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

[TopK] Follow-up fixes to the persistent top-k kernels

Browse Source 
- 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 <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 19:56:00 +07:00
co-authored by Hinata
parent 2918e3fda1
commit cff46af086
4 changed files with 215 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/persistent_topk.cuh
+136 -63
View File
@@ -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<int>(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<unsigned long long>(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<int>(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<uint32_t>(TopK)) {
if (slot < static_cast<uint32_t>(TopK) && !force_emit_fault) {
row_output[slot] = static_cast<int32_t>(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<unsigned long long>(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++;
}
csrc/libtorch_stable/topk.cu
+36 -5
View File
@@ -6,6 +6,7 @@
#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <mutex>
#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<uint32_t>(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<bool> 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<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);
"[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();
csrc/libtorch_stable/torch_bindings.cpp
+6 -2
View File
@@ -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
vllm/envs.py
+37
View File
@@ -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
),
}