diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 52a23ad089..8d601c2abc 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -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 __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 __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(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` 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 __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(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(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(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(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(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(TopK)) { row_output[slot] = static_cast(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(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(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( 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++; } } diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 0e6bee5253..19c7394ba9 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -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 g_exact_counters[kMaxCounterDevices] = {}; +std::atomic g_exact_counters[kMaxCounterDevices] = {}; std::atomic 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(P::RADIX); + 4u * ctas_per_group <= static_cast(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 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 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(); - unsigned int* src = + int64_t* dst = out.mutable_data_ptr(); + unsigned long long* src = ensure_exact_counters(static_cast(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"); diff --git a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py index 02c1720f11..08ffab8638 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py +++ b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py @@ -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))