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

[TopK] Deterministic placement and tie membership for the large-path radix_topk

Browse Source 
Above RADIX_THRESHOLD a row is split across a group of CTAs. The pivot was
already order-independent; cross-CTA placement, intra-CTA placement and tie
membership were arrival-order races. They are now deterministic behind
MERLIN_TOPK_DET_LARGE (default on).

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 19:55:59 +07:00
co-authored by Hinata
parent daa56356c8
commit 63e2d8be20
2 changed files with 384 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/persistent_topk.cuh
+323 -3
View File
@@ -147,6 +147,11 @@ struct PersistentTopKParams {
// nullptr; the kernel only ever atomicAdds into it, never reads back.
bool exact_medium;
unsigned int* exact_counters;
// MERLIN_TOPK_DET_LARGE, ANDed on the host with the capacity gate
// 2 * ctas_per_group <= RADIX (the per-CTA counts are staged in one 256-slot
// RadixRowState histogram row). Resolved once per process, so it is a
// launch constant and a cudagraph captures one value.
bool det_large;
};
// ============================================================================
@@ -468,7 +473,16 @@ enum TopKExactCounter : int {
kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop
kCtrMediumExactFallback = 2, // rows that took the exact restart
kCtrMediumExactFault = 3, // exact restart could not complete
kNumTopKExactCounters = 4,
// Large path (multi-CTA radix_topk, seq_len > RADIX_THRESHOLD).
// Slots 4-7 are only ever touched by the deterministic emission, so they
// stay 0 in the MERLIN_TOPK_DET_LARGE=0 control arm: `[9] == 1 && [4] == 0`
// while large rows are known to have run means the host-side capacity gate
// (2 * ctas_per_group <= RADIX) turned the switch off for that launch.
kCtrLargeDetRows = 4, // large rows emitted deterministically
kCtrLargeTiesCut = 5, // rows where the pivot-equal group was cut
kCtrLargeTiesDropped = 6, // pivot-equal elements discarded by those cuts
kCtrLargeFault = 7, // large-path invariant violation
kNumTopKExactCounters = 8,
};
// Exact, bounded-memory re-selection for a single medium-path row.
@@ -1074,6 +1088,300 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
// coordinated radix select via global-memory histograms and barriers.
// ============================================================================
// ============================================================================
// Deterministic placement and tie membership for the large path
//
// `radix_topk`'s pivot is already an exact, order-independent function of the
// row: every quantity feeding it is an integer histogram add and a
// single-threaded descending bin walk. Three things after the
// pivot are not:
//
// (a) cross-CTA placement — a CTA's base offset into `row_output` is its
// arrival order at `state->output_counter`;
// (b) intra-CTA placement — a slot inside that base is arrival order at a
// shared counter;
// (c) cross-CTA tie membership — which pivot-equal elements survive the
// `pos < TopK` cut is a race across CTAs.
//
// (a) and (b) change only the permutation of the selected list, but sparse
// attention consumes that list as an ordered gather, so a permutation forks
// the hidden state. (c) changes the set. All three are replaced here by a
// rank derived from (cta_in_group, element index):
//
// * each CTA counts its own `> pivot` and `== pivot` elements out of the
// chunk already resident in shared memory (no new global reads);
// * the counts are published into `RadixRowState::histogram[radix_iter % 3]`
// — the slot round 3 has just finished reading — so the workspace
// contract is unchanged;
// * every CTA reads all 2*ctas_per_group counts and derives its own base by
// an exclusive prefix over `cta_in_group`, in one fixed order, so all
// CTAs agree without further communication;
// * the tie cut keeps the HIGHEST global indices (most-recent-wins, the
// `MERLIN_TIE_RECENT` policy, measured at +3.4 % tokens/step
// over oldest-wins). Chunks are handed out in ascending index order
// (`my_chunk_start = cta_in_group * chunk_size`), so "the highest
// `equal_needed` pivot-equal indices" is exactly "walk CTAs from the top
// down", which the per-CTA equal counts already answer;
// * emission is a prefix sum over ascending element index — no atomic
// anywhere on the placement path — so the row is strictly ascending in
// index and bit-reproducible.
//
// Cost per row: one extra inter-CTA barrier, two extra ballots and one
// 32-lane shuffle scan per 1,024-element tile, 2 global stores and
// 2*ctas_per_group global loads. No second pass over the row in global
// memory: the pivot-equal count that the tie cut needs is produced by the same
// shared-memory pass v2.0.1 already runs to count `> pivot`.
// ============================================================================
template <int TopK>
__device__ __noinline__ void radix_emit_deterministic(
int32_t* __restrict__ row_output, uint32_t ordered_pivot,
uint32_t equal_needed, uint32_t my_chunk_start, uint32_t actual_chunk_size,
uint32_t* local_histogram, uint32_t* suffix_sum, uint32_t* shared_ordered,
RadixRowState* state, uint32_t cta_in_group, uint32_t ctas_per_group,
int& barrier_phase, uint32_t radix_iter, uint32_t tx,
unsigned int* counters) {
constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u;
// Staging for the published per-CTA counts. The host gates the switch on
// 2 * ctas_per_group <= RADIX, so these two halves never collide.
constexpr uint32_t kEqStage = RADIX / 2u;
// ---- per-CTA counts, out of the resident chunk -------------------------
if (tx < 2u) suffix_sum[tx] = 0u;
__syncthreads();
uint32_t my_gt = 0u;
uint32_t my_eq = 0u;
for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) {
const uint32_t ord = shared_ordered[i];
if (ord > ordered_pivot) {
my_gt++;
} else if (ord == ordered_pivot) {
my_eq++;
}
}
#pragma unroll
for (int off = 16; off > 0; off >>= 1) {
my_gt += __shfl_down_sync(0xffffffffu, my_gt, off);
my_eq += __shfl_down_sync(0xffffffffu, my_eq, off);
}
if ((tx & 31u) == 0u) {
// Integer counts: the atomic's arrival order cannot change the sum, so
// this is a reduction, not a placement.
if (my_gt > 0u) atomicAdd(&suffix_sum[0], my_gt);
if (my_eq > 0u) atomicAdd(&suffix_sum[1], my_eq);
}
__syncthreads();
const uint32_t local_gt = suffix_sum[0];
const uint32_t local_eq = suffix_sum[1];
// ---- barrier X1: the histogram slot is free ----------------------------
// Round 3's `current_hist` is state->histogram[(radix_iter*4+3) % 3], i.e.
// histogram[radix_iter % 3] (4r+3 == r mod 3). Every CTA copies it into
// shared memory AFTER round 3's own barrier and there is no inter-CTA
// barrier between that copy and here, so a fast CTA could otherwise
// overwrite the slot while a slow CTA is still reading it. This barrier is
// the one the slot-reuse argument requires.
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();
uint32_t* cta_counts = state->histogram[radix_iter % 3u];
if (tx == 0) {
cta_counts[cta_in_group] = local_gt;
cta_counts[ctas_per_group + cta_in_group] = local_eq;
}
// ---- barrier X2: the counts are published ------------------------------
// red_release issues `fence.acq_rel.gpu` before the arrival add and wait_ge
// reads with `ld.global.acquire.gpu`, so the two plain stores above are
// visible at device scope to every CTA that leaves this barrier. This is
// the same release/acquire pairing the shipped kernel already relies on for
// CTA-0's plain-store zeroing of `next_hist`.
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();
// ---- bases and the tie cut ---------------------------------------------
for (uint32_t c = tx; c < ctas_per_group; c += kThreadsPerBlock) {
local_histogram[c] = cta_counts[c];
local_histogram[kEqStage + c] = cta_counts[ctas_per_group + c];
}
__syncthreads();
if (tx == 0) {
uint32_t gt_base = 0u;
uint32_t gt_total = 0u;
for (uint32_t c = 0; c < ctas_per_group; ++c) {
const uint32_t g = local_histogram[c];
if (c < cta_in_group) gt_base += g;
gt_total += g;
}
// Most-recent-wins: walk CTAs from the top down, handing each one as many
// of the `equal_needed` slots as the CTAs above it left unclaimed. keep[]
// is written over eq[] in place (eq[c] is read before keep[c] is stored).
uint32_t eq_above = 0u;
for (int c = static_cast<int>(ctas_per_group) - 1; c >= 0; --c) {
const uint32_t e = local_histogram[kEqStage + c];
const uint32_t want =
(eq_above >= equal_needed) ? 0u : (equal_needed - eq_above);
local_histogram[kEqStage + c] = (want < e) ? want : e;
eq_above += e;
}
const uint32_t eq_total = eq_above;
uint32_t eq_keep_base = 0u;
for (uint32_t c = 0; c < cta_in_group; ++c) {
eq_keep_base += local_histogram[kEqStage + c];
}
const uint32_t my_keep = local_histogram[kEqStage + cta_in_group];
// Invariants, all three derived from the medium-path restart design and re-derivable
// here: the radix walk guarantees #{key > pivot} == TopK - equal_needed
// and equal_needed <= #{key == pivot}, and the two together make the
// emission a bijection onto [0, TopK). Every CTA evaluates this from the
// same published integers, so the verdict is group-uniform and needs no
// further communication.
uint32_t fault = 0u;
if (gt_total + equal_needed != static_cast<uint32_t>(TopK)) fault = 1u;
if (equal_needed > eq_total) fault = 1u;
if (my_keep > local_eq) fault = 1u;
suffix_sum[2] = gt_base;
suffix_sum[3] = eq_keep_base;
// Keep the local pivot-equal elements whose ascending rank within this
// chunk is >= cut_rank, i.e. the `my_keep` highest-indexed of them.
suffix_sum[4] = local_eq - my_keep;
suffix_sum[5] = fault;
if (counters != nullptr && cta_in_group == 0u) {
atomicAdd(&counters[kCtrLargeDetRows], 1u);
if (eq_total > equal_needed) {
atomicAdd(&counters[kCtrLargeTiesCut], 1u);
atomicAdd(&counters[kCtrLargeTiesDropped], eq_total - equal_needed);
}
if (fault != 0u) atomicAdd(&counters[kCtrLargeFault], 1u);
}
}
__syncthreads();
const uint32_t gt_base = suffix_sum[2];
const uint32_t eq_keep_base = suffix_sum[3];
const uint32_t cut_rank = suffix_sum[4];
const uint32_t row_fault = suffix_sum[5];
// ---- emission: rank by index, no atomics -------------------------------
if (row_fault == 0u) {
// Scratch for the tile scan. The staged per-CTA counts above are dead by
// now (thread 0 consumed them before the barrier), so local_histogram is
// free again; 33 slots each at 0 and at 64 stay inside RADIX.
uint32_t* warp_gt = local_histogram;
uint32_t* warp_eq = local_histogram + 64;
const uint32_t lane = tx & 31u;
const uint32_t warp = tx >> 5;
uint32_t gt_run = 0u;
uint32_t eq_run = 0u;
// Block-uniform trip count: every thread runs every tile, so the ballots
// and the barriers inside are reached by the whole block.
for (uint32_t base = 0u; base < actual_chunk_size;
base += kThreadsPerBlock) {
const uint32_t i = base + tx;
const bool in_range = (i < actual_chunk_size);
const uint32_t ord = in_range ? shared_ordered[i] : 0u;
const bool f_gt = in_range && (ord > ordered_pivot);
const bool f_eq = in_range && (ord == ordered_pivot);
const uint32_t bgt = __ballot_sync(0xffffffffu, f_gt);
const uint32_t beq = __ballot_sync(0xffffffffu, f_eq);
const uint32_t lane_mask = (1u << lane) - 1u;
const uint32_t gt_in_warp = static_cast<uint32_t>(__popc(bgt & lane_mask));
const uint32_t eq_in_warp = static_cast<uint32_t>(__popc(beq & lane_mask));
if (lane == 0u) {
warp_gt[warp] = static_cast<uint32_t>(__popc(bgt));
warp_eq[warp] = static_cast<uint32_t>(__popc(beq));
}
__syncthreads();
if (tx < 32u) {
const uint32_t cg = (tx < kWarpsPerBlock) ? warp_gt[tx] : 0u;
const uint32_t ce = (tx < kWarpsPerBlock) ? warp_eq[tx] : 0u;
uint32_t sg = cg;
uint32_t se = ce;
#pragma unroll
for (int off = 1; off < 32; off <<= 1) {
const uint32_t ng = __shfl_up_sync(0xffffffffu, sg, off);
const uint32_t ne = __shfl_up_sync(0xffffffffu, se, off);
if (tx >= static_cast<uint32_t>(off)) {
sg += ng;
se += ne;
}
}
warp_gt[tx] = sg - cg; // inclusive -> exclusive
warp_eq[tx] = se - ce;
if (tx == 31u) {
warp_gt[32] = sg; // tile totals (lanes >= kWarpsPerBlock add 0)
warp_eq[32] = se;
}
}
__syncthreads();
const uint32_t gt_rank = gt_run + warp_gt[warp] + gt_in_warp;
const uint32_t eq_rank = eq_run + warp_eq[warp] + eq_in_warp;
const uint32_t tile_gt = warp_gt[32];
const uint32_t tile_eq = warp_eq[32];
if (f_gt || (f_eq && eq_rank >= cut_rank)) {
// Selected elements before this one, in ascending index order:
// all `> pivot` ones, plus the pivot-equal ones past the cut.
const uint32_t kept_eq_before =
(eq_rank > cut_rank) ? (eq_rank - cut_rank) : 0u;
const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before;
if (slot < static_cast<uint32_t>(TopK)) {
row_output[slot] = static_cast<int32_t>(my_chunk_start + i);
} else if (counters != nullptr) {
// Unreachable while the three invariants above hold; bounded store
// plus a counter rather than an out-of-range write.
atomicAdd(&counters[kCtrLargeFault], 1u);
}
}
__syncthreads(); // before warp_gt / warp_eq are rewritten
gt_run += tile_gt;
eq_run += tile_eq;
}
// No thread may start the next row's chunk load — which writes
// shared_ordered — until every warp has finished reading it here.
__syncthreads();
} else {
// ---- fail loud ------------------------------------------------------
// The predicate is group-uniform (same published integers, same fixed
// summation order in every CTA), so no CTA of this group emitted anything
// and CTA 0 can fill the row alone. -1 is the kernel's own padding value
// and the consumer rejects it, so a faulted row attends nothing instead of
// quietly attending the wrong blocks.
if (cta_in_group == 0u) {
for (uint32_t i = tx; i < static_cast<uint32_t>(TopK);
i += kThreadsPerBlock) {
row_output[i] = -1;
}
}
__threadfence();
__syncthreads();
#if MERLIN_TOPK_EXACT_TRAP
__trap();
#endif
}
}
// ============================================================================
// Multi-CTA cooperative RadixTopK for a single large row.
// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
@@ -1087,7 +1395,8 @@ __device__ void radix_topk(const float* __restrict__ row_input,
uint32_t* shared_scalars, uint32_t* shared_ordered,
RadixRowState* state, uint32_t cta_in_group,
uint32_t ctas_per_group, int& barrier_phase,
uint32_t radix_iter, uint32_t tx) {
uint32_t radix_iter, uint32_t tx, bool det_large,
unsigned int* counters) {
const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
? my_chunk_start + chunk_size
: seq_len;
@@ -1241,6 +1550,16 @@ __device__ void radix_topk(const float* __restrict__ row_input,
// -- Count local > pivot elements --
const uint32_t ordered_pivot = shared_scalars[0];
// Deterministic placement + tie membership. Everything below this
// branch is v2.0.1 verbatim and is what MERLIN_TOPK_DET_LARGE=0 runs.
if (det_large) {
radix_emit_deterministic<TopK>(
row_output, ordered_pivot, shared_scalars[1], my_chunk_start,
actual_chunk_size, local_histogram, suffix_sum, shared_ordered, state,
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters);
return;
}
if (tx == 0) suffix_sum[0] = 0;
__syncthreads();
@@ -1389,7 +1708,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
radix_topk<TopK, VEC_SIZE>(
row_input, row_output, seq_len, my_chunk_start, chunk_size,
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx);
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx,
params.det_large, params.exact_counters);
radix_iter++;
}
}
csrc/libtorch_stable/topk.cu
+61 -7
View File
@@ -4,6 +4,7 @@
#include <cuda_runtime.h>
#include <algorithm>
#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <mutex>
@@ -37,6 +38,24 @@ bool exact_medium_enabled() {
return enabled;
}
// ---------------------------------------------------------------------------
// MERLIN_TOPK_DET_LARGE switch.
//
// Same contract as the medium-path switch: resolved once per process, so it is a
// launch constant and a cudagraph captures one value. Default ON on this
// branch. MERLIN_TOPK_DET_LARGE=0 is the A/B control arm and runs the
// v2.0.1 `radix_topk` tail verbatim -- the deterministic emission sits behind
// one `if (det_large)` in front of it and touches nothing when off, including
// the large-path counters, which is what makes the control arm the real baseline.
// ---------------------------------------------------------------------------
bool det_large_enabled() {
static const bool enabled = [] {
const char* raw = std::getenv("MERLIN_TOPK_DET_LARGE");
return (raw == nullptr || raw[0] == '\0') ? true : (std::atoi(raw) != 0);
}();
return enabled;
}
// Process-lifetime device counters, one block per device. The kernel only
// ever atomicAdds into them, so the launch path needs no host sync and the
// values accumulate across cudagraph replays.
@@ -356,6 +375,30 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
params.ctas_per_group = ctas_per_group;
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
params.exact_medium = exact_medium_enabled();
// The per-CTA selected / pivot-equal counts are staged in ONE
// 256-slot RadixRowState histogram row (the slot round 3 has just finished
// reading), which holds 2 * ctas_per_group uint32s. That is the whole
// reason the workspace contract is unchanged, and it is also the only
// limit: above 128 CTAs per group the deterministic path cannot be
// expressed there, so the launch degrades to the v2.0.1 tail rather than
// failing. Unreachable at any shape this kernel serves (ctas_per_group is
// 2-3 across our 135K-262K band and <= 16 at the smallest chunk size), but
// a silent loss of determinism is exactly what this branch exists to stop,
// so say it once.
const bool det_large_fits =
2u * ctas_per_group <= static_cast<uint32_t>(P::RADIX);
params.det_large = det_large_enabled() && det_large_fits;
if (det_large_enabled() && !det_large_fits) {
static std::atomic<bool> det_large_warned{false};
if (!det_large_warned.exchange(true)) {
std::fprintf(stderr,
"[vllm] MERLIN_TOPK_DET_LARGE is on but ctas_per_group=%u "
"needs %u of %d staging slots; the large path falls back "
"to nondeterministic placement at this shape (counter [4] "
"stays 0).\n",
ctas_per_group, 2u * ctas_per_group, P::RADIX);
}
}
// Allocate only off the capture path; in the ordinary engine lifecycle the
// eager warmup lands here first, long before capture. The explicit,
// quiescent-phase entry point is the reader op.
@@ -460,12 +503,17 @@ void persistent_topk(const torch::stable::Tensor& logits,
// counters. This is also the explicit pre-capture init path: call it once per
// device before any cudagraph capture.
//
// `out` is an int32 CUDA tensor with at least 5 elements:
// `out` is an int32 CUDA tensor with at least 10 elements:
// [0] medium-path rows entered (denominator)
// [1] rows where a 4,096-slot guard would drop a candidate
// [2] rows that took the exact restart (0 in the control arm)
// [3] exact restarts that could not complete (must stay 0)
// [4] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0
// [4] large-path rows emitted deterministically (0 in the control arm)
// [5] large-path rows whose pivot-equal group was cut at the k-th value
// [6] large-path pivot-equal elements discarded by those cuts
// [7] large-path invariant violations (must stay 0)
// [8] 1 if MERLIN_TOPK_EXACT_MEDIUM is on, else 0
// [9] 1 if MERLIN_TOPK_DET_LARGE is on, else 0
//
// COHERENCY -- stated precisely, because the earlier wording overclaimed.
// With `sync == 0` the snapshot is ordered ONLY against work already enqueued
@@ -500,9 +548,9 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset,
STD_TORCH_CHECK(out.is_contiguous(),
"out must be contiguous (expanded or strided views would be "
"written out of bounds or into the wrong elements)");
STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 1,
STD_TORCH_CHECK(out.numel() >= vllm::persistent::kNumTopKExactCounters + 2,
"out must hold at least ",
vllm::persistent::kNumTopKExactCounters + 1, " elements");
vllm::persistent::kNumTopKExactCounters + 2, " elements");
STD_TORCH_CHECK(
reset == 0 || sync != 0,
"persistent_topk_exact_counters: reset requires sync != 0, otherwise the "
@@ -547,10 +595,16 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset,
}
// Static storage: the source of an async H2D copy must outlive the copy.
static const int32_t enabled = exact_medium_enabled() ? 1 : 0;
// Two flags: [kNumTopKExactCounters] = the medium-path switch,
// [kNumTopKExactCounters + 1] = the large-path switch. The large
// flag reports the ENV value, not the per-launch effective value: a launch
// whose ctas_per_group overflowed the staging row still reports 1 here and
// leaves counter [4] at 0, which is how that degradation is diagnosed.
static const int32_t enabled[2] = {exact_medium_enabled() ? 1 : 0,
det_large_enabled() ? 1 : 0};
STD_TORCH_CHECK(
cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, &enabled,
sizeof(int32_t), cudaMemcpyHostToDevice,
cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, enabled,
2 * sizeof(int32_t), cudaMemcpyHostToDevice,
stream) == cudaSuccess,
"exact counter flag write failed");