Skip to content

[HIP] [Bugfix] sample_kernels: partial-vector and prologue-tail defects #7373

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
junhaha666 merged 1 commit into from
Nov 3, 2026
Merged

[HIP] [Bugfix] sample_kernels: partial-vector and prologue-tail defects #7373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 82 additions & 29 deletions csrc/kernels/sample_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,13 @@ __device__ void random_sample_outer_exponential_impl(const DTYPE_I* input,

int k = threadIdx.x * vec_size_i;
int vec_stride = BlockSize * vec_size_i;
// A thread whose first element is already past N never loads, so its
// registers stay undefined and `k` is never advanced; the tail below must
// not run for it (`k - vec_stride` would be negative and pass the check).
const bool has_pre = k < N;
vec_i vec_inp_pre;
vec_f vec_exp_pre;
if(k < N)
if(has_pre)
{
vec_inp_pre = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
vec_exp_pre = load_vector_nbytes<float, vec_size_i, load_chunk_bytes>(buffer_e, k);
Expand All @@ -65,11 +69,20 @@ __device__ void random_sample_outer_exponential_impl(const DTYPE_I* input,
temperature = 1.0f / temperature;

auto loop = [&]() {
// Elements this call covers are `base .. base + vec_size_i - 1`; the
// last vector of the row is partial when N % vec_size_i != 0. What a
// lane past N reads is not a fixed value -- `buffer_e` ends exactly at
// N so it yields 0, while `buffer_i`'s extent is rounded up to a dword
// (`oob_i`), so for a 2-byte dtype the first lane past an odd N still
// reads real memory. Neither must reach the running max or the argmax.
const int base = k - vec_stride;
vec_f vec_cur_f;
float new_max_softmax = max_softmax;
for(int i = 0; i < vec_size_i; i++)
{
vec_cur_f[i] = opus::cast<float>(vec_inp_pre[i]) * temperature;
vec_cur_f[i] = (base + i) < N
? opus::cast<float>(vec_inp_pre[i]) * temperature
: -FLT_MAX;
new_max_softmax = max(new_max_softmax, vec_cur_f[i]);
}
for(int i = 0; i < vec_size_i; i++)
Expand All @@ -93,9 +106,9 @@ __device__ void random_sample_outer_exponential_impl(const DTYPE_I* input,
{
vec_exp_pre[i] += eps;
vec_cur_f[i] = vec_cur_f[i] / vec_exp_pre[i];
if(vec_cur_f[i] > thread_kvp.value)
if((base + i) < N && vec_cur_f[i] > thread_kvp.value)
{
thread_kvp.key = k - vec_stride + i;
thread_kvp.key = base + i;
thread_kvp.value = vec_cur_f[i];
}
}
Expand All @@ -110,7 +123,7 @@ __device__ void random_sample_outer_exponential_impl(const DTYPE_I* input,
vec_exp_pre = vec_exp_cur;
}
// tail
if((k - vec_stride) < N)
if(has_pre && (k - vec_stride) < N)
{
loop();
}
Expand Down Expand Up @@ -191,8 +204,10 @@ __device__ void random_sample_impl(const DTYPE_I* input,

int k = threadIdx.x * vec_size_i;
int vec_stride = BlockSize * vec_size_i;
// See the matching comment in random_sample_outer_exponential_impl.
const bool has_pre = k < N;
vec_i vec_pre;
if(k < N)
if(has_pre)
{
vec_pre = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
k += vec_stride;
Expand All @@ -201,13 +216,16 @@ __device__ void random_sample_impl(const DTYPE_I* input,
temperature = 1.0f / temperature;

auto loop = [&]() {
auto rand = dist_func(&state);
vec_i vec_cur = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
const int base = k - vec_stride;
auto rand = dist_func(&state);
vec_i vec_cur = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
vec_f vec_cur_f;
float new_max_softmax = max_softmax;
for(int i = 0; i < vec_size_i; i++)
{
vec_cur_f[i] = opus::cast<float>(vec_pre[i]) * temperature;
vec_cur_f[i] = (base + i) < N
? opus::cast<float>(vec_pre[i]) * temperature
: -FLT_MAX;
new_max_softmax = max(new_max_softmax, vec_cur_f[i]);
}
for(int i = 0; i < vec_size_i; i++)
Expand All @@ -231,9 +249,9 @@ __device__ void random_sample_impl(const DTYPE_I* input,
{
float u = transform_func((&rand.x)[i]) + eps;
vec_cur_f[i] = vec_cur_f[i] / u;
if(vec_cur_f[i] > thread_kvp.value)
if((base + i) < N && vec_cur_f[i] > thread_kvp.value)
{
thread_kvp.key = k - vec_stride + i;
thread_kvp.key = base + i;
thread_kvp.value = vec_cur_f[i];
}
}
Expand All @@ -245,7 +263,7 @@ __device__ void random_sample_impl(const DTYPE_I* input,
vec_pre = loop();
}
// tail
if((k - vec_stride) < N)
if(has_pre && (k - vec_stride) < N)
{
loop();
}
Expand Down Expand Up @@ -296,37 +314,56 @@ __device__ void argmax_impl(const DTYPE_I* input, int* output, int m_idx, int N,

using kvp = aiter::KeyValuePair<int, float>;
aiter::ArgMax arg_max;
kvp thread_kvp{0, -FLT_MAX};
// -INFINITY, not -FLT_MAX: this compares raw logits, so -FLT_MAX is a value
// the input can actually take. A row whose only finite logit is -FLT_MAX
// would tie with the sentinel and lose it on ArgMax's smaller-key rule,
// returning token 0. Only a genuine -inf ties now, and token 0 is the
// documented answer for an all--inf row. The sampling impls keep -FLT_MAX:
// their candidates are exp(..)/u >= 0, so no real value can reach it, and
// -inf there would turn `thread_kvp.value * ratio` into NaN when the
// running-max rescale underflows ratio to 0.
kvp thread_kvp{0, -INFINITY};
int k = threadIdx.x * vec_size_i;
int vec_stride = BlockSize * vec_size_i;
// See the matching comment in random_sample_outer_exponential_impl.
const bool has_pre = k < N;
vec_i vec_pre;
if(k < N)
if(has_pre)
{
vec_pre = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
k += vec_stride;
}
for(; k < N; k += vec_stride)
{
kvp tmp_kvp{k - vec_stride - 1, -FLT_MAX};
vec_i vec_cur = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
// A lane past N reads whatever the bounded descriptor yields there, which
// wins the argmax whenever every real logit in the row is negative.
// Neutralise those lanes rather than shortening the loop: the trip count
// has to stay a compile-time constant or the vec_size_i-wide body stops
// being unrolled. Both halves of the pair are neutralised, so no key at or
// past N is ever handed to the reduction.
// The per-lane select costs greedy_sample ~3-4% at small M (16-wide body).
// Branching the whole-vector case around it -- only the row's last vector
// can straddle N -- was measured slower still (10.97 vs 10.36 us at M=1),
// so the straight-line select is what ships.
auto scan = [&](int base) {
kvp tmp_kvp{0, -INFINITY};
for(int i = 0; i < vec_size_i; i++)
{
tmp_kvp.key += 1;
tmp_kvp.value = opus::cast<float>(vec_pre[i]);
const bool in_range = (base + i) < N;
tmp_kvp.key = in_range ? base + i : 0;
tmp_kvp.value = in_range ? opus::cast<float>(vec_pre[i]) : -INFINITY;
thread_kvp = arg_max(thread_kvp, tmp_kvp);
}
};

for(; k < N; k += vec_stride)
{
vec_i vec_cur = load_vector_nbytes<DTYPE_I, vec_size_i, load_chunk_bytes>(buffer_i, k);
scan(k - vec_stride);
vec_pre = vec_cur;
}
// tail
if((k - vec_stride) < N)
if(has_pre && (k - vec_stride) < N)
{
kvp tmp_kvp{k - vec_stride - 1, -FLT_MAX};
for(int i = 0; i < vec_size_i; i++)
{
tmp_kvp.key += 1;
tmp_kvp.value = opus::cast<float>(vec_pre[i]);
thread_kvp = arg_max(thread_kvp, tmp_kvp);
}
scan(k - vec_stride);
}

thread_kvp = block_reduce<kvp, aiter::ArgMax, BlockSize, true>(thread_kvp, arg_max);
Expand Down Expand Up @@ -721,7 +758,23 @@ __global__ void exponential_kernel(DTYPE_O* output,
float u = transform_func((&rand.x)[i]) + eps;
vec_cur[i] = opus::cast<DTYPE_O>(u);
}
store_vector<DTYPE_O, DTYPE_O, vec_size_o, RT, false>(buffer_o, vec_cur, k, 1.0f);
if(k + vec_size_o <= N)
{
store_vector<DTYPE_O, DTYPE_O, vec_size_o, RT, false>(buffer_o, vec_cur, k, 1.0f);
}
else
{
// Partial last vector. The descriptor cannot mask it: `oob_o`
// rounds the extent up to a dword because the buffer bound is
// enforced per dword, so at a 2-byte dtype element N-1 and element
// N share one and the hardware can only keep or drop both. Keeping
// both writes past the row; dropping both loses a real element.
// Only one thread per row ever lands here, once.
for(int i = 0; k + i < N; i++)
{
ptr_o[k + i] = vec_cur[i];
}
}
}
}

Expand Down
165 changes: 165 additions & 0 deletions op_tests/test_sample.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,165 @@ def test_mixed_sample(M, N, dtype=torch.bfloat16, eps=1e-6):
}


def check_exponential_row_bounds(l_dtype=(torch.float32, torch.bfloat16)):
"""aiter.exponential must fill its row and write nothing past it.

Same partial-last-vector shape as check_vocab_edges, on the store side.
The output descriptor's extent is rounded up to a dword (`oob_o`) because
the buffer bound is enforced per dword, so at a 2-byte dtype with
N % VecSize != 0 the last store reached one element past the row and a
two-row buffer had row 1 clobbered by row 0's noise. Shrinking the extent
to N instead drops a real element, since at an odd N elements N-1 and N
share a dword -- hence the masked scalar tail in the kernel.
"""
sentinel = -7.0
for dtype in l_dtype:
for N in (4094, 4095, 4096, 4097, 16385):
buf = torch.full((2, N), sentinel, dtype=dtype)
aiter.exponential(buf[0:1], lambd=1.0, eps=1e-10)
torch.cuda.synchronize()
unwritten = int((buf[0] == sentinel).sum())
assert (
unwritten == 0
), f"exponential {dtype} {N=}: {unwritten} elements left unwritten"
assert bool(
(buf[1] == sentinel).all()
), f"exponential {dtype} {N=}: wrote past the row into the next one"
aiter.logger.info(
"exponential row-bounds check passed for %s", [str(d) for d in l_dtype]
)


def check_vocab_edges(l_dtype=(torch.bfloat16, torch.float32), eps=1e-10):
"""Every sampler must return a token id in [0, N) for any vocab size.

Two shapes used to escape that, and neither is reachable from the perf
sweep's N (both of its values are multiples of the vector width and far
above one block's span):

* N <= (BlockSize - 1) * VecSize: threads whose first element is past N
skipped the prologue load, then ran the tail on undefined registers
with a negative base index.
* N % VecSize != 0: the row's last vector is partial, and the lanes past
N fall outside the row. For the outer-exponential kernels the noise
descriptor ends exactly at N, so the 0 it yields lands in the divisor
and exp(x - max) / (0 + eps) won the argmax for any logits. The other
kernels take their out-of-range value from the logits descriptor
instead, so it only wins when every real logit in the row is negative
-- hence the all-negative rows below.

Both block spans are covered: 16384 for greedy_sample (block 1024, 16-wide)
and 4096 for every other entry point (block 1024, 4-wide), including
argmax_impl as reached through the mixed kernels at temperature 0.

fp32 is swept as well as bf16: `oob_i` rounds the logits descriptor up to a
dword, so at a 2-byte dtype the first lane past an odd N still reads real
memory and the corruption is data-dependent, while at fp32 it is a hard 0.
"""
l_vocab = [
256,
1000,
2048,
4092,
4095,
4096,
4097,
4099,
8193,
16368,
16383,
16385,
50001,
163840,
163841,
]
M = 4
for dtype in l_dtype:
for N in l_vocab:
for all_negative in (False, True):
logits = torch.randn(M, N, device="cuda", dtype=torch.float32) * 4
if all_negative:
logits = -logits.abs() - 1.0
logits = logits.to(dtype)
ref_argmax = logits.to(torch.float).argmax(dim=-1).to(torch.int32)
hot = torch.ones(M, dtype=torch.float32, device="cuda")
cold = torch.zeros(M, dtype=torch.float32, device="cuda")
exponential = torch.empty(M, N, dtype=torch.float32).exponential_(1)
tag = f"{dtype} {N=} {all_negative=}"

out = torch.empty(M, dtype=torch.int32, device="cuda")
aiter.greedy_sample(out, logits)
assert torch.equal(out, ref_argmax), f"greedy_sample {tag} {out=}"

for name, fn in (
(
"mixed_sample_outer_exponential",
aiter.mixed_sample_outer_exponential,
),
(
"random_sample_outer_exponential",
aiter.random_sample_outer_exponential,
),
):
out = torch.empty(M, dtype=torch.int32, device="cuda")
fn(out, logits, exponential, hot, eps=eps)
probs = softmax(logits.to(torch.float))
ref = probs.div_(exponential + eps).argmax(dim=-1).to(torch.int32)
assert torch.equal(out, ref), f"{name} {tag} {out=}"

# temperature 0 routes the mixed kernels into argmax_impl, which
# the greedy entry point only ever reaches 16 elements wide.
out = torch.empty(M, dtype=torch.int32, device="cuda")
aiter.mixed_sample_outer_exponential(
out, logits, exponential, cold, eps=eps
)
assert torch.equal(
out, ref_argmax
), f"mixed_sample_outer_exponential T=0 {tag} {out=}"

out = torch.empty(M, dtype=torch.int32, device="cuda")
aiter.mixed_sample(out, logits, cold, lambd=1.0, eps=eps)
assert torch.equal(out, ref_argmax), f"mixed_sample T=0 {tag} {out=}"

# The inner-exponential kernels draw their own noise, so only
# the index is checkable against a reference.
for name, fn in (
("mixed_sample", aiter.mixed_sample),
("random_sample", aiter.random_sample),
):
out = torch.empty(M, dtype=torch.int32, device="cuda")
fn(out, logits, hot, lambd=1.0, eps=eps)
assert bool(
((out >= 0) & (out < N)).all()
), f"{name} {tag} out of range: {out=}"
# The most negative finite value of the dtype is something logits can
# actually take, so it must not tie with argmax_impl's sentinel and lose it
# on the smaller-key rule. Taken per dtype: -FLT_MAX is outside bf16's
# range and rounds to -inf, the one case where token 0 is the right answer.
for dtype in l_dtype:
neg_max = -torch.finfo(dtype).max
for N, tok in ((17, 5), (4095, 4094), (16385, 16384)):
logits = torch.full((1, N), float("-inf"), dtype=torch.float32)
logits[0, tok] = neg_max
logits = logits.to(dtype)
assert torch.isfinite(logits[0, tok]), f"{dtype} cannot hold {neg_max}"
cold = torch.zeros(1, dtype=torch.float32, device="cuda")
e = torch.empty(1, N, dtype=torch.float32).exponential_(1)
tag = f"{dtype} {N=} finite-logit=-FLT_MAX at {tok}"

out = torch.empty(1, dtype=torch.int32, device="cuda")
aiter.greedy_sample(out, logits)
assert int(out[0]) == tok, f"greedy_sample {tag} -> {int(out[0])}"

out = torch.empty(1, dtype=torch.int32, device="cuda")
aiter.mixed_sample_outer_exponential(out, logits, e, cold, eps=eps)
assert int(out[0]) == tok, f"mixed T=0 {tag} -> {int(out[0])}"

aiter.logger.info(
"vocab-edge check passed over %s for %s", l_vocab, [str(d) for d in l_dtype]
)


d_sample = {
"greedy": test_greedy_sample,
"random": test_random_sample,
Expand Down Expand Up @@ -242,3 +401,9 @@ def test_mixed_sample(M, N, dtype=torch.bfloat16, eps=1e-6):
df = pd.DataFrame(df)
df_md = df.to_markdown(index=False)
aiter.logger.info("sample summary (markdown):\n%s", df_md)

# After the sweep: this asserts, and it draws from the same global RNG the
# sweep does, so running it first would both hide the perf tables on a
# regressed tree and change the inputs every table above is measured on.
check_vocab_edges()
check_exponential_row_bounds()
Loading