From cff46af086528749b48a8df1a9cf8a6556d2dd26 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 23 Sep 2026 08:35:50 -0400 Subject: [PATCH] [TopK] Follow-up fixes to the persistent top-k kernels - Counter enum declared above the decode template (two-phase lookup). - Emission-fault flag set atomically. - Close the legacy large tie loop with a block barrier (inherited). - Large-path fault injectors registered as hashed passthroughs. - The tensor-free counter-slots op registered for CompositeExplicitAutograd. - Tie measurements recorded only on rows neither fault verdict failed. - Tests scoped to the properties the build promises. Co-Authored-By: Hinata Signed-off-by: MOVZX --- csrc/libtorch_stable/persistent_topk.cuh | 199 ++++++++++++++++------- csrc/libtorch_stable/topk.cu | 41 ++++- csrc/libtorch_stable/torch_bindings.cpp | 8 +- vllm/envs.py | 37 +++++ 4 files changed, 215 insertions(+), 70 deletions(-) diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index f02aebfc16..504db92a92 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -189,6 +189,87 @@ struct PersistentTopKParams { // 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; + // Which fault the injector forces in that CTA. + // false (MERLIN_TOPK_LARGE_FAULT_KIND unset or `pre`): the PRE-EMISSION + // predicate `my_keep > local_eq`, caught at X3. Emission is skipped + // entirely on that row, so this kind cannot reach the X4 code at all. + // true (`emit`): the EMISSION-BOUND slot check inside the placement loop, + // caught at X4. This is the only way to exercise `suffix_sum[6]`, and + // it exercises it the way the hardware would: every emitting thread of + // the named CTA discovers an invalid slot at once, which is exactly the + // concurrency the atomic exists for. + // Resolved once per process on the host, so it is a launch constant. + bool fault_inject_emit; +}; + +// ============================================================================ +// Device-global exact-path counters -- DECLARED HERE, ABOVE EVERY USER. +// +// This enum used to live beside the medium-path restart, +// ~200 lines BELOW the decode template that increments +// `kCtrDecodeOverflow`. `counters` is a plain `unsigned long long*` +// parameter, so that enumerator is a NON-DEPENDENT name: two-phase lookup +// resolves it when the template is DEFINED, not when it is instantiated, +// and the Linux build rejected the tree with +// persistent_topk.cuh(442): identifier "kCtrDecodeOverflow" is undefined. +// An offline `nvcc -cubin` on a non-Linux host does NOT reproduce that on its own: nvcc +// can drive cicc in Microsoft mode (`--ms_c++17 --microsoft_version=`), and that mode +// defers non-dependent lookup to instantiation, where the enum is already +// visible. `-Xcudafe --dep_name` restores strict two-phase lookup and the +// offline harness then reproduces the Linux diagnostic verbatim. +// +// The fix is the declaration ORDER, not a literal: every increment still +// names its enumerator, so a slot renumbering stays a one-line edit here. +// ============================================================================ + +// Device-global counters. Plain device atomics on a process-lifetime buffer: +// 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. +// +// SLOT MAP (the tie-determinism + large-path integration, 11 slots). The tie change numbered its decode +// counter AROUND the large-path block rather than after the medium-path one, so the merged layout is +// the union of the two branches with nothing renumbered and no external probe +// silently re-reading a row count as a switch -- the failure a merged layout +// must avoid. +// +// [0] medium rows [1] medium overflow [2] medium exact fallback +// [3] medium exact fault +// [4]-[7] large path (rows_large, ties_cut, ties_dropped, faults_large) +// [8] MERLIN_TOPK_EXACT_MEDIUM [9] MERLIN_TOPK_DET_LARGE +// [10] decode overflow +// +// Slots 8-9 are written by the HOST after the device copy; the device never +// touches them, and they are slots INSIDE the array, not two elements appended +// after it. A stale five- or ten-element reader buffer is rejected outright by +// the length check in the reader op, which is the fail-loud direction. +enum TopKExactCounter : int { + kCtrMediumRows = 0, // denominator: medium-path rows entered + 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 + // 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 + // (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 + kCtrLargeFault = 7, // large-path invariant violation + kSlotSwitchMedium = 8, // host-written flag, NOT a device counter + kSlotSwitchLarge = 9, // host-written flag, NOT a device counter + kCtrDecodeOverflow = 10, // decode rows whose buffered-candidate count + // exceeded DBUF before deterministic ranking + kNumTopKExactCounters = 11, // total slots: device buffer AND reader output }; // ============================================================================ @@ -604,56 +685,6 @@ __device__ __noinline__ bool histogram_2048_topk( #define MERLIN_TOPK_EXACT_TRAP 1 #endif -// Device-global counters. Plain device atomics on a process-lifetime buffer: -// 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. -// -// SLOT MAP (the tie-determinism + large-path integration, 11 slots). The tie change numbered its decode -// counter AROUND the large-path block rather than after the medium-path one, so the merged layout is -// the union of the two branches with nothing renumbered and no external probe -// silently re-reading a row count as a switch -- the failure a merged layout -// must avoid. -// -// [0] medium rows [1] medium overflow [2] medium exact fallback -// [3] medium exact fault -// [4]-[7] large path (rows_large, ties_cut, ties_dropped, faults_large) -// [8] MERLIN_TOPK_EXACT_MEDIUM [9] MERLIN_TOPK_DET_LARGE -// [10] decode overflow -// -// Slots 8-9 are written by the HOST after the device copy; the device never -// touches them, and they are slots INSIDE the array, not two elements appended -// after it. A stale five- or ten-element reader buffer is rejected outright by -// the length check in the reader op, which is the fail-loud direction. -enum TopKExactCounter : int { - kCtrMediumRows = 0, // denominator: medium-path rows entered - 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 - // 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 - // (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 - kCtrLargeFault = 7, // large-path invariant violation - kSlotSwitchMedium = 8, // host-written flag, NOT a device counter - kSlotSwitchLarge = 9, // host-written flag, NOT a device counter - kCtrDecodeOverflow = 10, // decode rows whose buffered-candidate count - // exceeded DBUF before deterministic ranking - kNumTopKExactCounters = 11, // total slots: device buffer AND reader output -}; - // Exact, bounded-memory re-selection for a single medium-path row. // // `scores` is already offset to the row. Every radix round rescans the whole @@ -1419,7 +1450,8 @@ __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 long long* counters, int fault_inject_cta, bool tie_recent) { + unsigned long long* counters, int fault_inject_cta, + bool fault_inject_emit, bool tie_recent) { constexpr uint32_t kWarpsPerBlock = kThreadsPerBlock / 32u; // Staging for the published per-CTA counts. The host gates the switch on // 4 * ctas_per_group <= RADIX, so these regions never collide. @@ -1585,7 +1617,7 @@ __device__ __noinline__ void radix_emit_deterministic( // 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 && + if (fault_inject_cta >= 0 && !fault_inject_emit && static_cast(cta_in_group) == fault_inject_cta) { my_keep_checked = local_eq + 1u; } @@ -1622,22 +1654,22 @@ __device__ __noinline__ void radix_emit_deterministic( barrier_phase++; __syncthreads(); + // Tie measurement staged in thread 0's register; recorded only after BOTH + // verdicts (X3 here, X4 after emission) have cleared the row. Recording it + // here, behind X3 alone, counted tie cuts for rows an emission-bound fault + // later failed closed (reference host, FAULT_KIND=emit: ties_cut +8 on 8 + // all -1 rows). + uint32_t staged_ties_dropped = 0u; 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; + staged_ties_dropped = suffix_sum[8]; if (counters != nullptr && cta_in_group == 0u) { 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])); - } } } __syncthreads(); @@ -1648,6 +1680,17 @@ __device__ __noinline__ void radix_emit_deterministic( const uint32_t keep_width = suffix_sum[9]; const uint32_t row_fault = suffix_sum[5]; + // TEST ONLY: force the emission-bound violation in one + // named CTA. Block-uniform (both operands are launch constants and + // `cta_in_group` is block-uniform), so it changes no control flow the whole + // block does not take together. It deliberately fires for EVERY emitting + // thread of that CTA rather than one: the defect under repair is a + // concurrent write to `suffix_sum[6]`, and a single-writer injection would + // pass with the broken code. + const bool force_emit_fault = + fault_inject_emit && fault_inject_cta >= 0 && + static_cast(cta_in_group) == fault_inject_cta; + // ---- emission: rank by index, no atomics ------------------------------- if (row_fault == 0u) { // Scratch for the tile scan. The staged per-CTA counts above are dead by @@ -1724,7 +1767,7 @@ __device__ __noinline__ void radix_emit_deterministic( (eq_rank > cut_rank) ? (eq_rank - cut_rank) : 0u; if (kept_eq_before > keep_width) kept_eq_before = keep_width; const uint32_t slot = gt_base + eq_keep_base + gt_rank + kept_eq_before; - if (slot < static_cast(TopK)) { + if (slot < static_cast(TopK) && !force_emit_fault) { row_output[slot] = static_cast(my_chunk_start + i); } else { // An out-of-range slot used to bump a counter and @@ -1732,7 +1775,16 @@ __device__ __noinline__ void radix_emit_deterministic( // 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; + // + // ATOMIC, not a plain store. Every thread whose + // slot is out of range reaches this line, and an invariant failure + // is precisely the situation in which several of them do at once -- + // so the plain `suffix_sum[6] = 1u` was a concurrent non-atomic + // write inside the mechanism that exists to detect the failure. It + // is a shared-memory word and the flag is monotonic, so `atomicOr` + // is the whole fix; the healthy path never executes it. The medium + // restart already uses atomicOr for the same job. + atomicOr(&suffix_sum[6], 1u); } } @@ -1765,6 +1817,14 @@ __device__ __noinline__ void radix_emit_deterministic( emit_fault |= cta_counts[kEmitStage + c]; } suffix_sum[7] = emit_fault; + // 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 (counters != nullptr && cta_in_group == 0u && row_fault == 0u && + emit_fault == 0u && staged_ties_dropped > 0u) { + atomicAdd(&counters[kCtrLargeTiesCut], 1ull); + atomicAdd(&counters[kCtrLargeTiesDropped], + static_cast(staged_ties_dropped)); + } } __syncthreads(); @@ -1814,7 +1874,7 @@ __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 tie_recent) { + bool fault_inject_emit, bool tie_recent) { const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) ? my_chunk_start + chunk_size : seq_len; @@ -1975,7 +2035,7 @@ __device__ void radix_topk(const float* __restrict__ row_input, 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, - fault_inject_cta, tie_recent); + fault_inject_cta, fault_inject_emit, tie_recent); return; } @@ -2029,6 +2089,18 @@ __device__ void radix_topk(const float* __restrict__ row_input, } } } + // BASELINE repair (inherited from v2.0.1): the tie loop + // above reads shared_ordered, and the next row this CTA processes -- a + // large row's Stage-1 chunk load, or a decode/medium row reusing the same + // dynamic shared memory -- writes it before its first barrier. Without a + // closing barrier a fast warp can overwrite entries a slow warp is still + // comparing against the pivot, changing tie MEMBERSHIP, not only order. + // Every thread of the block reaches this point (the loop bound is + // block-uniform and there is no early return in the legacy tail), so the + // barrier is uniform. Applies to both comparison arms: the legacy path is + // the MERLIN_TOPK_DET_LARGE=0 control, so the control now carries this + // repair too, like the histogram-publication barrier (an earlier commit). + __syncthreads(); } // ============================================================================ @@ -2142,6 +2214,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2) 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.fault_inject_cta, + params.fault_inject_emit, effective_tie_recent(params.det_ties, params.tie_recent)); radix_iter++; } diff --git a/csrc/libtorch_stable/topk.cu b/csrc/libtorch_stable/topk.cu index 53b2ee0718..fed5fc885a 100644 --- a/csrc/libtorch_stable/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include "ops.h" @@ -132,6 +133,31 @@ int large_fault_inject_cta() { return cta; } +// --------------------------------------------------------------------------- +// MERLIN_TOPK_LARGE_FAULT_KIND. TEST ONLY, companion to +// the variable above; ignored unless MERLIN_TOPK_LARGE_FAULT_INJECT names a +// CTA. +// +// unset / "pre" (default) force the PRE-EMISSION predicate, caught at X3. +// "emit" force the EMISSION-BOUND slot check, caught at X4. +// +// The two are mutually exclusive by construction: a pre-emission fault makes +// `row_fault` nonzero, the emission block is skipped wholesale, and the X4 +// code -- including `suffix_sum[6]`, the word whose concurrent non-atomic +// write this change is about -- is never reached. The shipped injector could +// therefore only ever exercise the arm that was already correct, which is why +// the X4 arm had never been exercised. Resolved once per process like +// every other switch here, so it is a launch constant a cudagraph bakes in. +// --------------------------------------------------------------------------- +bool large_fault_inject_emit() { + static const bool emit = [] { + const char* raw = std::getenv("MERLIN_TOPK_LARGE_FAULT_KIND"); + if (raw == nullptr || raw[0] == '\0') return false; + return std::strncmp(raw, "emit", 5) == 0; + }(); + return emit; +} + // --------------------------------------------------------------------------- // Say the effective tie policy out loud, once. // @@ -511,6 +537,7 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, 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(); + params.fault_inject_emit = large_fault_inject_emit(); if (det_large_enabled() && !det_large_fits) { static std::atomic det_large_warned{false}; if (!det_large_warned.exchange(true)) { @@ -526,11 +553,15 @@ void launch_persistent_topk(const torch::stable::Tensor& logits, 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); + "[vllm] MERLIN_TOPK_LARGE_FAULT_INJECT=%d kind=%s is " + "set: the large path FORCES a %s 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, + params.fault_inject_emit ? "emit" : "pre", + params.fault_inject_emit ? "emission-bound (X4)" + : "pre-emission (X3)"); } } params.det_ties = det_ties_enabled(); diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 3a9613329c..344aa96faf 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -865,8 +865,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); ops.impl("persistent_topk_exact_counters", TORCH_BOX(&persistent_topk_exact_counters)); - ops.impl("persistent_topk_exact_counters_slots", - TORCH_BOX(&persistent_topk_exact_counters_slots)); #ifdef VLLM_ENABLE_COOPERATIVE_TOPK ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk)); #endif @@ -943,6 +941,12 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) { ops.impl("cutlass_scaled_mm_supports_fp4", TORCH_BOX(&cutlass_scaled_mm_supports_fp4)); #endif + // Tensor-free: the counter layout the BUILD carries. Registered here, not + // in the CUDA block -- with no tensor argument there is no dispatch key, so + // a CUDA-only registration can never be called (reference host: + // NotImplementedError "no fallback function is registered"). + ops.impl("persistent_topk_exact_counters_slots", + TORCH_BOX(&persistent_topk_exact_counters_slots)); } // Cache ops diff --git a/vllm/envs.py b/vllm/envs.py index 6f18de6837..88b17ac4bc 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -568,6 +568,35 @@ if TYPE_CHECKING: # bool that would imply this tree can act on it. MERLIN_TOPK_DET_LARGE: str | None = None # + # The large-path FAULT INJECTORS. TEST + # ONLY, and registered for exactly the same reason as the switches above, + # not because anything should set them in production: + # + # * both are resolved once per process by a std::getenv in topk.cu + # (large_fault_inject_cta(), large_fault_inject_emit()) and the + # resolved values reach the GPU as kernel launch arguments, so a + # cudagraph BAKES THEM IN. An injected build and a clean build differ + # in the captured graph while sharing a cache key -- which is the + # failure mode: a graph captured under injection returns an all -1 row + # for every deterministic large row, and on replay nothing re-reads the + # environment to notice the variable has since been unset. FAULT_INJECT + # was once "unregistered and capture-baked"; hashing + # them is what stops a cached artifact crossing between the arms. + # * compile_factors() is a deny-list, so a name it does not know cannot + # be hashed at all. Neither is in ignored_factors. + # + # Raw passthroughs, like the switches: the C++ getters are the single + # source of truth (INJECT parses with atoi and treats <0 and unset alike; + # KIND compares against the literal "emit"). Do not parse them here. + # + # Which CTA of the group has its fault predicate forced; unset / negative + # is off, and setting it also suppresses __trap(). + MERLIN_TOPK_LARGE_FAULT_INJECT: str | None = None + # Which fault: unset / "pre" forces the pre-emission predicate caught at + # X3, "emit" forces the emission-bound slot check caught at X4. Ignored + # unless MERLIN_TOPK_LARGE_FAULT_INJECT names a CTA. + MERLIN_TOPK_LARGE_FAULT_KIND: str | None = None + # # 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 @@ -2645,6 +2674,14 @@ environment_variables: dict[str, Callable[[], Any]] = { # MERLIN_TOPK_EXACT_MEDIUM is registered once, earlier in this dict. "MERLIN_TOPK_EXACT_DECODE": lambda: os.getenv("MERLIN_TOPK_EXACT_DECODE", None), "MERLIN_TOPK_DET_LARGE": lambda: os.getenv("MERLIN_TOPK_DET_LARGE", None), + # Large-path fault injectors. TEST ONLY, hashed for + # the reason in the declaration block: both are capture-baked. + "MERLIN_TOPK_LARGE_FAULT_INJECT": lambda: os.getenv( + "MERLIN_TOPK_LARGE_FAULT_INJECT", None + ), + "MERLIN_TOPK_LARGE_FAULT_KIND": lambda: os.getenv( + "MERLIN_TOPK_LARGE_FAULT_KIND", None + ), }