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

[TopK] Large-path hardening

Browse Source 
- Inline-PTX acquire/release helpers declare their memory effects.
- The fault verdict is aggregated across the CTA group before emission, so a
  faulted row is failed closed by every CTA.
- 64-bit device counters and an int64 reader (counter protocol v2).
- Tests: adversarial distributions, slot-residue coverage, an sm_90 gate for
  tests that need it.

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 63e2d8be20
commit 512ca2fc0b
3 changed files with 314 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/persistent_topk.cuh
+212 -53
View File
@@ -146,12 +146,19 @@ struct PersistentTopKParams {
// is a process-lifetime device buffer of kNumTopKExactCounters uint32s, or
// nullptr; the kernel only ever atomicAdds into it, never reads back.
bool exact_medium;
unsigned int* exact_counters;
unsigned long long* 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;
// TEST ONLY. -1 (the default) disables it; any
// other value names the `cta_in_group` whose CTA-LOCAL fault predicate is
// forced, so a fault originating outside CTA 0 can be exercised. Resolved
// once per process from MERLIN_TOPK_LARGE_FAULT_INJECT, so it is a launch
// constant like the two switches. Setting it also suppresses __trap(), because
// the point of injecting is to observe the trap-off row-failure policy.
int fault_inject_cta;
};
// ============================================================================
@@ -468,6 +475,14 @@ __device__ __noinline__ void histogram_2048_topk(
// no host sync and no host-visible allocation on the launch path, so this is
// safe to record inside a cudagraph capture and on every replay. Read from
// Python with `torch.ops._C.persistent_topk_exact_counters`.
//
// The counters are `unsigned long long`, not `unsigned int`,
// and the reader tensor is int64. A u32 counter copied into a SIGNED int32
// tensor goes negative, then wraps: at 65,536 columns with k=512 an all-equal
// row drops 65,024 ties, so slot 6 crosses the signed boundary after 33,027
// rows and wraps after 66,053, and a delta taken across that boundary is a
// wrong number that looks like a measurement. `atomicAdd(unsigned long long*)`
// is native on sm_86, so the width costs nothing on the increment path.
enum TopKExactCounter : int {
kCtrMediumRows = 0, // denominator: medium-path rows entered
kCtrMediumOverflow = 1, // rows where a 4,096-slot guard would drop
@@ -477,7 +492,7 @@ enum TopKExactCounter : int {
// 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.
// (4 * 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
@@ -494,7 +509,7 @@ enum TopKExactCounter : int {
template <int TopK>
__device__ __noinline__ void medium_exact_restart(
const float* __restrict__ scores, int32_t* __restrict__ output,
uint32_t length, unsigned int* __restrict__ counters) {
uint32_t length, unsigned long long* __restrict__ counters) {
static_assert(kThreadsPerBlock >= RADIX,
"the single-thread bin walk assumes RADIX <= blockDim.x");
@@ -759,7 +774,7 @@ __device__ __noinline__ void medium_exact_restart(
// ---- fail loud ---------------------------------------------------------
if (row_fault != 0u) {
if (tx == 0 && counters != nullptr) {
atomicAdd(&counters[kCtrMediumExactFault], 1u);
atomicAdd(&counters[kCtrMediumExactFault], 1ull);
}
// -1 is the kernel's own padding value and the consumer rejects it
// (`valid = ... & (token >= 0)`), so a faulted row attends nothing and is
@@ -790,7 +805,7 @@ template <int TopK>
__device__ __noinline__ void histogram_256_topk(
const float* __restrict__ logits, int* __restrict__ output_indices,
int logits_offset, int seq_len, bool exact_medium,
unsigned int* __restrict__ counters) {
unsigned long long* __restrict__ counters) {
// No default arguments on purpose: a new call site must state whether it
// wants the exact path, rather than silently inheriting the broken one.
// All shared state lives in dynamic shared memory to avoid static
@@ -815,7 +830,7 @@ __device__ __noinline__ void histogram_256_topk(
if (thread_id == 0) {
shared_overflow_counted = 0;
if (counters != nullptr) atomicAdd(&counters[kCtrMediumRows], 1u);
if (counters != nullptr) atomicAdd(&counters[kCtrMediumRows], 1ull);
}
if (thread_id < RADIX + 1) {
@@ -891,7 +906,7 @@ __device__ __noinline__ void histogram_256_topk(
// Past the early return the threshold bin WILL be collected, so the M1
// prediction is now a real overflow event.
if (thread_id == 0 && shared_overflow_counted == 1 && counters != nullptr) {
atomicAdd(&counters[kCtrMediumOverflow], 1u);
atomicAdd(&counters[kCtrMediumOverflow], 1ull);
}
__syncthreads();
@@ -951,11 +966,11 @@ __device__ __noinline__ void histogram_256_topk(
if (__builtin_expect(raw_buffered > MAX_BUFFERED_ITEMS, 0)) {
if (thread_id == 0 && shared_overflow_counted == 0) {
shared_overflow_counted = 1;
if (counters != nullptr) atomicAdd(&counters[kCtrMediumOverflow], 1u);
if (counters != nullptr) atomicAdd(&counters[kCtrMediumOverflow], 1ull);
}
if (exact_medium) {
if (thread_id == 0 && counters != nullptr) {
atomicAdd(&counters[kCtrMediumExactFallback], 1u);
atomicAdd(&counters[kCtrMediumExactFallback], 1ull);
}
__syncthreads(); // no thread may clobber medium_smem before this
medium_exact_restart<TopK>(logits + logits_offset, output_indices,
@@ -1036,24 +1051,55 @@ __device__ __noinline__ void histogram_256_topk(
// Inter-CTA sync primitives
// ============================================================================
// `asm volatile` alone is NOT a compiler memory barrier. It
// stops the statement being deleted or duplicated and keeps volatile asms in
// order relative to one another, but it says nothing about the surrounding
// loads and stores, so nvcc was free to sink a publication store past the
// release fence or hoist a payload load above the acquire load. NVIDIA
// documents exactly this in "Inline PTX Assembly: Incorrect Optimization".
// Every helper below therefore declares a `"memory"` clobber, and the fence is
// fused into the SAME asm statement as the atomic it orders so that no
// optimizer pass can separate them.
//
// Why inline PTX and not `cuda::atomic_ref` / `__threadfence()`:
// * `__threadfence()` compiles to `membar.gl` on this toolchain (verified:
// nvcc 12.8 -arch=sm_86 -ptx), which is the sequentially-consistent device
// fence, strictly stronger than the `fence.acq_rel.gpu` this protocol
// needs. Swapping it in would change the emitted ordering primitive on
// the hottest barrier in the kernel, which is a performance change smuggled
// into a correctness fix.
// * `cuda::atomic_ref<int, thread_scope_device>` expresses the same
// acquire/release semantics, but it would pull libcu++ into a header that
// is compiled for every arch this extension builds, and `red` (a reduction
// with no return value) has no atomic_ref spelling -- `fetch_add` returns
// the old value and costs the return path.
// So: keep the PTX, and make it honest about its memory effects.
__device__ __forceinline__ int ld_acquire(int* ptr) {
int state = 0;
#if (__CUDA_ARCH__ >= 700)
asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"
: "=r"(state)
: "l"(ptr));
: "l"(ptr)
: "memory");
#else
asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr));
asm volatile("ld.cg.global.b32 %0, [%1];\n"
: "=r"(state)
: "l"(ptr)
: "memory");
#endif
return state;
}
__device__ __forceinline__ void red_release(int* ptr, int val) {
#if (__CUDA_ARCH__ >= 700)
asm volatile("fence.acq_rel.gpu;\n");
asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"
:
: "l"(ptr), "r"(val));
// One statement: the fence and the arrival add cannot be separated, and the
// clobber forbids moving any access across either of them.
asm volatile(
"fence.acq_rel.gpu;\n"
"red.relaxed.gpu.global.add.s32 [%0], %1;\n"
:
: "l"(ptr), "r"(val)
: "memory");
#else
__threadfence();
atomicAdd(ptr, val);
@@ -1062,8 +1108,12 @@ __device__ __forceinline__ void red_release(int* ptr, int val) {
__device__ __forceinline__ void st_release(int* ptr, int val) {
#if (__CUDA_ARCH__ >= 700)
asm volatile("fence.acq_rel.gpu;\n");
asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val));
asm volatile(
"fence.acq_rel.gpu;\n"
"st.release.gpu.global.b32 [%0], %1;\n"
:
: "l"(ptr), "r"(val)
: "memory");
#else
__threadfence();
atomicExch(ptr, val);
@@ -1126,11 +1176,20 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
// 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
// * fault handling is fail-closed for the WHOLE GROUP: each CTA publishes a
// verdict word and every CTA ORs all of them (X3) before any CTA emits,
// and the emission bound is re-aggregated the same way (X4), so a fault
// raised in CTA 1 still produces an all-`-1` row and exactly one counter
// increment.
//
// Cost per row: FOUR inter-CTA barriers where the v2.0.1 tail runs one (X1
// slot reuse, X2 counts, X3 fault verdict, X4 emission verdict), two extra
// ballots and one 32-lane shuffle scan per 1,024-element tile, 4 global stores
// and 4*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`.
// shared-memory pass v2.0.1 already runs to count `> pivot`. X3/X4 are the
// price of a group-uniform fail-closed contract; at ctas_per_group 2-3 each is
// one release add and a short spin.
// ============================================================================
template <int TopK>
__device__ __noinline__ void radix_emit_deterministic(
@@ -1139,14 +1198,27 @@ __device__ __noinline__ void radix_emit_deterministic(
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) {
unsigned long long* counters, int fault_inject_cta) {
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.
// 4 * ctas_per_group <= RADIX, so these regions never collide.
constexpr uint32_t kEqStage = RADIX / 2u;
// Two more regions of the same published histogram row
// carry the per-CTA fault verdicts -- the pre-emission one at [2C, 3C) and
// the emission-bound one at [3C, 4C). They are SEPARATE regions on purpose:
// reusing [2C, 3C) for the second round would let a fast CTA overwrite a word
// a slow CTA had not yet read. Each CTA writes only its own word, so nothing
// has to pre-zero them. Host gate: 4 * ctas_per_group <= RADIX.
const uint32_t kFaultStage = 2u * ctas_per_group;
const uint32_t kEmitStage = 3u * ctas_per_group;
// suffix_sum scratch map (shared, RADIX words):
// [0] gt count [1] eq count [2] gt_base [3] eq_keep_base [4] cut_rank
// [5] group fault verdict [6] this CTA's emission-bound violation
// [7] group emission-bound verdict [8] ties dropped on this row
// ---- per-CTA counts, out of the resident chunk -------------------------
if (tx < 2u) suffix_sum[tx] = 0u;
if (tx == 6u) suffix_sum[6] = 0u;
__syncthreads();
uint32_t my_gt = 0u;
@@ -1200,7 +1272,11 @@ __device__ __noinline__ void radix_emit_deterministic(
// 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`.
// CTA-0's plain-store zeroing of `next_hist`. Caveat: that
// argument is about the GPU memory model and only holds if the COMPILER also
// emits the stores before the fence; the `"memory"` clobber on red_release
// is what makes that true, and the PTX for this site is the evidence quoted
// in the commit message.
if (tx == 0) {
red_release(&state->arrival_counter, 1);
}
@@ -1244,31 +1320,75 @@ __device__ __noinline__ void radix_emit_deterministic(
}
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.
// Invariants. 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).
//
// The first two ARE group-uniform -- every CTA derives
// them from the same published integers in the same fixed order -- but the
// third, `my_keep > local_eq`, compares a value that came back through
// global memory against this CTA's own shared-memory count, and a local
// inconsistency is exactly what it exists to detect. It can therefore be
// true in CTA 1 and false in CTA 0, and the previous code let CTA 0 emit a
// perfectly ordinary-looking row while CTA 1 silently took the fault arm.
// The verdict is now published and ORed across the group below (X3) before
// ANY CTA emits. The fourth check is new and bounds this CTA's highest
// emission slot ahead of time, so the per-element `slot < TopK` test in the
// loop is a backstop rather than the only guard.
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;
uint32_t my_keep_checked = my_keep;
// Test-only: force the CTA-LOCAL predicate in one named CTA, so a fault
// that originates outside CTA 0 can be observed. Off (-1) in production;
// the host resolves it once per process, so it is a launch constant.
if (fault_inject_cta >= 0 &&
static_cast<int>(cta_in_group) == fault_inject_cta) {
my_keep_checked = local_eq + 1u;
}
if (my_keep_checked > local_eq) fault = 1u;
if (gt_base + eq_keep_base + local_gt + my_keep >
static_cast<uint32_t>(TopK)) {
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;
suffix_sum[4] = (my_keep <= local_eq) ? (local_eq - my_keep) : 0u;
suffix_sum[8] = (eq_total > equal_needed) ? (eq_total - equal_needed) : 0u;
// Publish this CTA's verdict. Own word, so no pre-zeroing is needed.
cta_counts[kFaultStage + cta_in_group] = fault;
}
// ---- barrier X3: the fault verdicts are published ----------------------
// Unconditional, so the per-CTA barrier count stays group-uniform.
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 (tx == 0) {
uint32_t row_fault = 0u;
for (uint32_t c = 0; c < ctas_per_group; ++c) {
row_fault |= cta_counts[kFaultStage + c];
}
suffix_sum[5] = row_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);
atomicAdd(&counters[kCtrLargeDetRows], 1ull);
// Only meaningful on a healthy row: on a faulted one the tie arithmetic
// is what is under suspicion, so do not record it as a measurement.
if (row_fault == 0u && suffix_sum[8] > 0u) {
atomicAdd(&counters[kCtrLargeTiesCut], 1ull);
atomicAdd(&counters[kCtrLargeTiesDropped],
static_cast<unsigned long long>(suffix_sum[8]));
}
if (fault != 0u) atomicAdd(&counters[kCtrLargeFault], 1u);
}
}
__syncthreads();
@@ -1347,10 +1467,13 @@ __device__ __noinline__ void radix_emit_deterministic(
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);
} else {
// An out-of-range slot used to bump a counter and
// carry on, so with trapping ON it neither trapped nor produced the
// promised all-`-1` row. It is now recorded as a row failure and
// goes through the SAME policy as the pre-emission faults below: no
// store here, and the row is filled with -1 after the group agrees.
suffix_sum[6] = 1u;
}
}
@@ -1361,23 +1484,58 @@ __device__ __noinline__ void radix_emit_deterministic(
// 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.
}
// ---- barrier X4: the emission verdicts are published -------------------
// Unconditional (both arms of the branch above reach it), so the per-CTA
// barrier count stays group-uniform. On the healthy path this costs one
// inter-CTA barrier per large row and nothing else: every word written here
// is this CTA's own.
if (tx == 0) {
cta_counts[kEmitStage + cta_in_group] = suffix_sum[6];
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 (tx == 0) {
uint32_t emit_fault = 0u;
for (uint32_t c = 0; c < ctas_per_group; ++c) {
emit_fault |= cta_counts[kEmitStage + c];
}
suffix_sum[7] = emit_fault;
}
__syncthreads();
// ---- fail closed, for the whole group ----------------------------------
// Either verdict is now group-uniform, so every CTA agrees on whether this
// row failed, and CTA 0 can overwrite the row alone -- including the case
// where some other CTA had already emitted before an out-of-range slot was
// discovered. -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. The counter is incremented exactly ONCE per failed row.
if (row_fault != 0u || suffix_sum[7] != 0u) {
if (cta_in_group == 0u) {
for (uint32_t i = tx; i < static_cast<uint32_t>(TopK);
i += kThreadsPerBlock) {
row_output[i] = -1;
}
if (tx == 0 && counters != nullptr) {
atomicAdd(&counters[kCtrLargeFault], 1ull);
}
}
__threadfence();
__syncthreads();
#if MERLIN_TOPK_EXACT_TRAP
__trap();
// Fault injection is a test mode and deliberately suppresses the trap:
// the point of injecting is to OBSERVE the trap-off policy (all-`-1` row,
// counter += 1), which a context-killing __trap() would make unobservable
// without a second build.
if (fault_inject_cta < 0) {
__trap();
}
#endif
}
}
@@ -1396,7 +1554,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
RadixRowState* state, uint32_t cta_in_group,
uint32_t ctas_per_group, int& barrier_phase,
uint32_t radix_iter, uint32_t tx, bool det_large,
unsigned int* counters) {
unsigned long long* counters, int fault_inject_cta) {
const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
? my_chunk_start + chunk_size
: seq_len;
@@ -1556,7 +1714,8 @@ __device__ void radix_topk(const float* __restrict__ row_input,
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);
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx, counters,
fault_inject_cta);
return;
}
@@ -1709,7 +1868,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
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,
params.det_large, params.exact_counters);
params.det_large, params.exact_counters, params.fault_inject_cta);
radix_iter++;
}
}
csrc/libtorch_stable/topk.cu
+88 -22
View File
@@ -56,6 +56,30 @@ bool det_large_enabled() {
return enabled;
}
// ---------------------------------------------------------------------------
// MERLIN_TOPK_LARGE_FAULT_INJECT. TEST ONLY.
//
// -1 (the default, and the value of an unset or negative variable) is off.
// Any other value names the `cta_in_group` whose CTA-LOCAL fault predicate the
// large path forces, which is the only way to reach the fail-closed arm: no
// legal input can make a CTA's own pivot-equal count disagree with the one it
// published, so those branches are otherwise untestable (as with the
// medium-path restart's own fault path). Resolved once per process, so
// it is a launch constant and a cudagraph captures one value. Setting it also
// suppresses __trap(), because the point of injecting is to OBSERVE the
// trap-off row-failure policy (all -1 row, counter += 1) without a second
// build.
// ---------------------------------------------------------------------------
int large_fault_inject_cta() {
static const int cta = [] {
const char* raw = std::getenv("MERLIN_TOPK_LARGE_FAULT_INJECT");
if (raw == nullptr || raw[0] == '\0') return -1;
const int v = std::atoi(raw);
return (v < 0) ? -1 : v;
}();
return cta;
}
// 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.
@@ -86,12 +110,12 @@ enum CounterState : int {
// a second host thread may be publishing them from the slow path. Release on
// publish, acquire on read, and the pointer is always published BEFORE the
// state that advertises it.
std::atomic<unsigned int*> g_exact_counters[kMaxCounterDevices] = {};
std::atomic<unsigned long long*> g_exact_counters[kMaxCounterDevices] = {};
std::atomic<int> g_exact_counters_state[kMaxCounterDevices] = {};
std::mutex g_exact_counters_mutex;
// Lookup only. Never allocates, never blocks, safe from any context.
unsigned int* get_exact_counters(int device_index) {
unsigned long long* get_exact_counters(int device_index) {
if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr;
if (g_exact_counters_state[device_index].load(std::memory_order_acquire) !=
kCountersReady) {
@@ -106,7 +130,7 @@ unsigned int* get_exact_counters(int device_index) {
// `stream_is_capturing` only proves the CALLING stream is idle of capture, so
// the intended entry point is the reader op called explicitly before capture;
// the non-capturing launch path calls this as a fallback.
unsigned int* ensure_exact_counters(int device_index, cudaStream_t stream) {
unsigned long long* ensure_exact_counters(int device_index, cudaStream_t stream) {
if (device_index < 0 || device_index >= kMaxCounterDevices) return nullptr;
const int state =
@@ -130,8 +154,8 @@ unsigned int* ensure_exact_counters(int device_index, cudaStream_t stream) {
}
constexpr size_t bytes =
vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int);
unsigned int* ptr = nullptr;
vllm::persistent::kNumTopKExactCounters * sizeof(unsigned long long);
unsigned long long* ptr = nullptr;
if (cudaMalloc(&ptr, bytes) == cudaSuccess && ptr != nullptr) {
// Host acquire/release orders the HOST state word; it says
// nothing about the device. cudaMemset on device memory is asynchronous
@@ -377,17 +401,20 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
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
// reading). Two more per-CTA regions were added for the
// group-uniform fault verdicts, so the row now holds 4 * ctas_per_group
// uint32s. That is still the whole reason the workspace contract is
// unchanged, and it is also the only
// limit: above 64 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);
4u * ctas_per_group <= static_cast<uint32_t>(P::RADIX);
params.det_large = det_large_enabled() && det_large_fits;
params.fault_inject_cta = large_fault_inject_cta();
if (det_large_enabled() && !det_large_fits) {
static std::atomic<bool> det_large_warned{false};
if (!det_large_warned.exchange(true)) {
@@ -396,7 +423,18 @@ void launch_persistent_topk(const torch::stable::Tensor& logits,
"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);
ctas_per_group, 4u * ctas_per_group, P::RADIX);
}
}
if (params.fault_inject_cta >= 0) {
static std::atomic<bool> inject_warned{false};
if (!inject_warned.exchange(true)) {
std::fprintf(stderr,
"[vllm] MERLIN_TOPK_LARGE_FAULT_INJECT=%d is set: the "
"large path FORCES a fault in that CTA of every "
"deterministic large row and returns an all -1 row. This "
"is a test mode and must not be set in production.\n",
params.fault_inject_cta);
}
}
// Allocate only off the capture path; in the ordinary engine lifecycle the
@@ -503,7 +541,31 @@ 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 10 elements:
// PROTOCOL VERSION 2 -- BREAKING. `out` is an **int64**
// CUDA tensor with at least 10 elements; int32 is REJECTED, not silently
// accepted. The device counters are 64-bit unsigned and are copied D2D
// without conversion, so the reader's element width must match the device's.
//
// Why the break rather than a reinterpreting reader: the counters were u32
// copied into a SIGNED int32 tensor. At 65,536 columns with k=512 an
// all-equal row drops 65,024 ties, so [6] goes negative in Python after 33,027
// rows and wraps modulo 2^32 after 66,053 -- and a subtraction across the
// signed boundary gives a wrong delta, which is worse than a missing number
// because it looks like a measurement. Keeping u32 and reinterpreting in the
// reader would need an enforced per-window increment bound plus quiescent
// resets, i.e. a contract every future caller has to remember. 64-bit device
// counters make the failure unreachable instead: `atomicAdd(unsigned long
// long*)` is native on sm_86 (it is a single ATOM.E.ADD.S64, not a CAS loop),
// so the instrument costs the same, and a 64-bit slot cannot be driven to the
// signed boundary by any run this kernel will ever see.
//
// Everything that reads these counters must move together: the Python reader
// `vllm/models/qwen4_exp/nvidia/ops/topk_counters.py` now allocates int64, and
// the medium-path tests allocate int64. An int32 caller gets a loud error rather
// than half a number. The OPERATOR SCHEMA is unchanged (dtype is not part of
// it), so a stale caller fails at the check, not at the binding.
//
// Slots:
// [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)
@@ -537,13 +599,17 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset,
int64_t sync) {
#ifndef USE_ROCM
STD_TORCH_CHECK(out.is_cuda(), "out must be a CUDA tensor");
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Int,
"out must be int32");
STD_TORCH_CHECK(
out.scalar_type() == torch::headeronly::ScalarType::Long,
"out must be int64: the exact-restart counters are 64-bit as of the "
"large-path counter protocol change, and an int32 buffer would be filled "
"from half of each counter. Allocate torch.zeros(10, dtype=torch.int64)."
" (int32 was protocol version 1 and is deliberately rejected.)");
// Element count alone is not enough. A CUDA scalar expanded
// to five elements (stride 0) satisfies numel but owns one element's
// storage, and the five consecutive raw int32 writes below would run past
// it; an ordinary strided view would receive the values in the wrong slots.
// Require a contiguous 1-D buffer before taking the pointer.
// to ten elements (stride 0) satisfies numel but owns one element's storage,
// and the raw int64 writes below would run past it; an ordinary strided view
// would receive the values in the wrong slots. Require a contiguous 1-D
// buffer before taking the pointer.
STD_TORCH_CHECK(out.dim() == 1, "out must be 1-D");
STD_TORCH_CHECK(out.is_contiguous(),
"out must be contiguous (expanded or strided views would be "
@@ -563,10 +629,10 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset,
"persistent_topk_exact_counters must not be called during "
"cudagraph capture");
constexpr size_t bytes =
vllm::persistent::kNumTopKExactCounters * sizeof(unsigned int);
vllm::persistent::kNumTopKExactCounters * sizeof(unsigned long long);
int32_t* dst = out.mutable_data_ptr<int32_t>();
unsigned int* src =
int64_t* dst = out.mutable_data_ptr<int64_t>();
unsigned long long* src =
ensure_exact_counters(static_cast<int>(out.get_device_index()), stream);
// An unavailable buffer used to be reported as four zeros
@@ -600,11 +666,11 @@ void persistent_topk_exact_counters(torch::stable::Tensor& out, int64_t reset,
// 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,
static const int64_t enabled[2] = {exact_medium_enabled() ? 1 : 0,
det_large_enabled() ? 1 : 0};
STD_TORCH_CHECK(
cudaMemcpyAsync(dst + vllm::persistent::kNumTopKExactCounters, enabled,
2 * sizeof(int32_t), cudaMemcpyHostToDevice,
2 * sizeof(int64_t), cudaMemcpyHostToDevice,
stream) == cudaSuccess,
"exact counter flag write failed");
vllm/models/qwen4_exp/nvidia/ops/topk_counters.py
+14 -4
View File
@@ -2,12 +2,20 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Python reader for the medium-path exact-restart top-k counters.
The counters themselves are device-global uint32s owned by
The counters themselves are device-global uint64s owned by
`csrc/libtorch_stable/topk.cu`; `torch.ops._C.persistent_topk_exact_counters`
snapshots them into an int32 CUDA tensor. Without this module nothing in a
snapshots them into an **int64** CUDA tensor. Without this module nothing in a
running engine ever calls that op, so the instrument records into a buffer no
one reads -- which measures nothing.
DTYPE IS PART OF THE PROTOCOL (large-path determinism change). Version 1 was u32
device counters read into a signed int32 tensor, which goes negative after
~33k all-equal rows at k=512 and wraps after ~66k, so a delta taken across the
boundary is a wrong number that looks like a measurement. Version 2 widens
both sides, and the op REJECTS an int32 buffer rather than filling it from half
of each counter. An older build whose op still demands int32 will therefore
raise here; that is the intended signal, not a regression to paper over.
SLOT LAYOUT is build-dependent, so nothing here hard-codes five. `slot_count()`
asks the build first (a `..._slots` op, should a later build expose one) and
otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5); names come from a
@@ -84,7 +92,7 @@ def _get_op(name: str):
def slot_count() -> int:
"""How many int32 slots this build's counter array carries."""
"""How many int64 slots this build's counter array carries."""
query = _get_op(_OP_NAME + "_slots")
if query is not None:
return int(query())
@@ -118,7 +126,9 @@ def read(reset: bool = False, sync: bool = True) -> dict[str, int]:
return {}
count = slot_count()
device = torch.device("cuda", torch.cuda.current_device())
out = torch.zeros(count, dtype=torch.int32, device=device)
# int64: the device counters are 64-bit and the op refuses any other
# width. See the module docstring -- this is protocol version 2.
out = torch.zeros(count, dtype=torch.int64, device=device)
op(out, 1 if reset else 0, 1 if sync else 0)
return dict(zip(slot_names(count), out.tolist(), strict=True))