From ce42d99b71a467d5c0066f33c2e73e55123364b3 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 13:47:34 -0400 Subject: [PATCH 1/7] tests/rocm: reproduce the uninitialised last_valid_id read Fails on the current kernels. A row with no positive probability satisfies no predicate, so nothing is ever marked valid and the fallback reads `temp_storage.last_valid_id`, which ROCm declares but never initialises: include/flashinfer/rocm/sampling.cuh:675 if (tx == 0 && max_valid_index != -1) { temp_storage->last_valid_id = max_valid_index; // only write Upstream initialises it to -1 at five sites; the ROCm header has none, so the read returns whatever the previous block left in LDS: E AssertionError: out-of-range token id: tensor([ 656060220, -1077873287, -1066197537, 847098865, ...]) top_k_sampling_from_probs does not even get that far -- it feeds the garbage straight into `probs[row_idx * d + sampled_id]` and aborts the process. The assertion is on range and `valid`, not on a sentinel: there is no -1 to match today. The batch is 256 rows so LDS is recycled across blocks; a single block can read a plausible in-range value and hide it. min_p is deliberately absent -- its predicate is `x >= pivot` with pivot = max_val * p, so an all-zero row makes every element qualify and the fallback is never taken. --- tests/rocm/test_sampling.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/rocm/test_sampling.py b/tests/rocm/test_sampling.py index a4db832bd7..9134e863d3 100644 --- a/tests/rocm/test_sampling.py +++ b/tests/rocm/test_sampling.py @@ -658,6 +658,36 @@ def test_a_tensor_seed_is_rejected_rather_than_collapsed(op): ) +@pytest.mark.parametrize( + "op, kwargs", + [ + ("sampling_from_probs", {}), + ("top_p_sampling_from_probs", {"top_p": 0.9}), + ("top_k_sampling_from_probs", {"top_k": 10}), + ("top_k_top_p_sampling_from_probs", {"top_k": 10, "top_p": 0.9}), + ], +) +def test_a_row_with_no_positive_probability_reports_invalid(op, kwargs): + """No element satisfies the predicate, so no index is ever marked valid. + + `last_valid_id` is never initialised on ROCm, so the fallback reads whatever + the previous block left in LDS -- asserting on a sentinel would test the + wrong thing. The batch is wide enough to recycle LDS across blocks; a single + block can read a plausible in-range value and hide the bug. + """ + batch, vocab = 256, 512 + probs = torch.zeros(batch, vocab, device="cuda") + + samples, valid = getattr(flashinfer.sampling, op)( + probs, **kwargs, return_valid=True + ) + + assert torch.all((samples >= 0) & (samples < vocab)), ( + f"out-of-range token id: {samples[(samples < 0) | (samples >= vocab)][:8]}" + ) + assert not bool(valid.any()), "a row with no positive probability is not valid" + + def test_every_row_reports_valid(): """ROCm's kernels have no reject path, so return_valid is uniformly true. From edd871ce35baafccf6ea41171560dd3ad928fb1b Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 14:24:24 -0400 Subject: [PATCH 2/7] rocm/sampling: initialise last_valid_id, and report `valid` per row Arms the sentinel at the five sites upstream arms it (three inside the retry loops, so they re-arm per round), then guards the fallback read the way upstream does: output 0 and valid=false rather than using the sentinel as an index. The initialisation is the load-bearing half. Without it `last_valid_id` is declared and only ever written under `if (max_valid_index != -1)`, so the fallback read returned whatever the previous block left in LDS -- a guard on == -1 alone would never have fired. Before this, the repro produced token ids like 656060220 and -1077873287, and top_k_sampling_from_probs fed one straight into `probs[row_idx * d + sampled_id]` and aborted the process. Five kernels and five launchers gain `bool* valid`, and `mark_all_valid` is replaced by `check_valid_out`, which keeps the shape checks it was also doing. The launchers pass arguments through type-erased `void* args[]` arrays that no compiler checks against the kernel signature, so each kernel, its launcher, its args array and its C++ call site were changed and built one at a time rather than in a batch. chain_speculative_sampling is untouched: upstream has no `valid` output there and never initialises the sentinel either, so both platforms read uninitialised LDS. Inherited deliberately -- there is nowhere in that API to report it. --- csrc/rocm/sampling.cu | 25 +++--- include/flashinfer/rocm/sampling.cuh | 116 +++++++++++++++++++-------- 2 files changed, 98 insertions(+), 43 deletions(-) diff --git a/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 052748fbf1..581ebddcfa 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -38,13 +38,13 @@ inline void reject_per_request_seed(const std::optional& maybe_seed_ "pass int seed/offset instead"); } -// The ROCm kernels carry the last-valid-index fallback but not upstream's reject -// flag, so every row yields a sample and `valid` is uniformly true. -inline void mark_all_valid(at::Tensor valid, unsigned int batch_size) { +// The kernels write `valid` per row now; this only checks the buffer they write +// into. It replaced a fill_(true), which is why the shape checks live here. +inline void check_valid_out(at::Tensor valid, unsigned int batch_size) { CHECK_INPUT(valid); CHECK_DIM(1, valid); CHECK_EQ(valid.size(0), static_cast(batch_size)); - valid.fill_(true); + TORCH_CHECK(valid.scalar_type() == at::kBool, "valid must be bool, got ", valid.scalar_type()); } void softmax(at::Tensor workspace_buffer, at::Tensor logits, at::Tensor output, @@ -106,12 +106,13 @@ void sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor valid, unsigned int vocab_size = probs.size(1); reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::SamplingFromProb( static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, vocab_size, deterministic, static_cast(philox_seed), static_cast(philox_offset), stream); @@ -133,12 +134,13 @@ void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v bool has_top_p_arr = maybe_top_p_arr.has_value(); reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::TopPSamplingFromProb( static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, top_p_val, vocab_size, deterministic, static_cast(philox_seed), @@ -164,12 +166,13 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v bool has_top_k_arr = maybe_top_k_arr.has_value(); reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::TopKSamplingFromProb( static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, top_k_val, vocab_size, deterministic, static_cast(philox_seed), @@ -195,14 +198,14 @@ void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v bool has_min_p_arr = maybe_min_p_arr.has_value(); reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::MinPSamplingFromProb( static_cast(probs.data_ptr()), has_min_p_arr ? static_cast(maybe_min_p_arr->data_ptr()) : nullptr, - static_cast(output.data_ptr()), + static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, min_p_val, vocab_size, deterministic, static_cast(philox_seed), static_cast(philox_offset), stream); @@ -230,7 +233,7 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te bool has_top_p_arr = maybe_top_p_arr.has_value(); reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); @@ -238,7 +241,7 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te static_cast(probs.data_ptr()), has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, - static_cast(output.data_ptr()), + static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, top_k_val, top_p_val, vocab_size, deterministic, static_cast(philox_seed), static_cast(philox_offset), stream); diff --git a/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index 2cebe7c784..69c6c9abfb 100644 --- a/include/flashinfer/rocm/sampling.cuh +++ b/include/flashinfer/rocm/sampling.cuh @@ -775,8 +775,8 @@ __global__ void SamplingFromLogitsKernel(DType* logits, IdType* output, IdType* template -__global__ void SamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { +__global__ void SamplingFromProbKernel(DType* probs, IdType* output, bool* valid, IdType* indices, + uint32_t d, uint64_t philox_seed, uint64_t philox_offset) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; hiprand_init(philox_seed, bx, philox_offset, &state); @@ -789,6 +789,7 @@ __global__ void SamplingFromProbKernel(DType* probs, IdType* output, IdType* ind reinterpret_cast&>( smem_sampling); temp_storage.sampled_id = d; + temp_storage.last_valid_id = -1; __syncthreads(); vec_t probs_vec; @@ -814,17 +815,28 @@ __global__ void SamplingFromProbKernel(DType* probs, IdType* output, IdType* ind // NOTE(Zihao): this would happen when u is very close to 1 // and the sum of probabilities is smaller than u // In this case, we use the last valid index as the sampled id + if (temp_storage.last_valid_id == -1) { + if (tx == 0) { + output[bx] = 0; + valid[bx] = false; + } + return; + } sampled_id = temp_storage.last_valid_id; } - output[bx] = sampled_id; + if (tx == 0) { + output[bx] = sampled_id; + valid[bx] = true; + } } template -__global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, - IdType* top_k_arr, uint32_t top_k_val, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { +__global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, bool* valid, + IdType* indices, IdType* top_k_arr, uint32_t top_k_val, + uint32_t d, uint64_t philox_seed, + uint64_t philox_offset) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; @@ -848,6 +860,7 @@ __global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, IdType* do { round += 1; temp_storage.sampled_id = d; + temp_storage.last_valid_id = -1; __syncthreads(); float u = hiprand_uniform(&state) * q; aggregate = 0; @@ -871,6 +884,13 @@ __global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, IdType* // NOTE(Zihao): this would happen when u is very close to 1 // and the sum of probabilities is smaller than u // In this case, we use the last valid index as the sampled id + if (temp_storage.last_valid_id == -1) { + if (tx == 0) { + output[bx] = 0; + valid[bx] = false; + } + return; + } sampled_id = temp_storage.last_valid_id; } double pivot_0 = probs[row_idx * d + sampled_id]; @@ -933,15 +953,17 @@ __global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, IdType* __syncthreads(); if (tx == 0) { output[bx] = sampled_id; + valid[bx] = true; } } template -__global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* indices, - float* top_p_arr, float top_p_val, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { +__global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, bool* valid, + IdType* indices, float* top_p_arr, float top_p_val, + uint32_t d, uint64_t philox_seed, + uint64_t philox_offset) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; @@ -963,6 +985,7 @@ __global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* int sampled_id; do { temp_storage.sampled_id = d; + temp_storage.last_valid_id = -1; __syncthreads(); float u = hiprand_uniform(&state) * q; aggregate = 0; @@ -986,6 +1009,13 @@ __global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* // NOTE(Zihao): this would happen when u is very close to 1 // and the sum of probabilities is smaller than u // In this case, we use the last valid index as the sampled id + if (temp_storage.last_valid_id == -1) { + if (tx == 0) { + output[bx] = 0; + valid[bx] = false; + } + return; + } sampled_id = temp_storage.last_valid_id; } double pivot_0 = probs[row_idx * d + sampled_id]; @@ -1044,6 +1074,7 @@ __global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* __syncthreads(); if (tx == 0) { output[bx] = sampled_id; + valid[bx] = true; } } @@ -1051,8 +1082,9 @@ template __global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdType* output, - IdType* indices, float min_p_val, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { + bool* valid, IdType* indices, float min_p_val, + uint32_t d, uint64_t philox_seed, + uint64_t philox_offset) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; float p = (min_p_arr == nullptr) ? min_p_val : min_p_arr[bx]; hiprandStatePhilox4_32_10_t state; @@ -1101,6 +1133,7 @@ __global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdTyp int sampled_id; temp_storage.sampled_id = d; + temp_storage.last_valid_id = -1; __syncthreads(); float u = hiprand_uniform(&state) * q; #pragma unroll 2 @@ -1122,18 +1155,28 @@ __global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdTyp // NOTE(Zihao): this would happen when u is very close to 1 // and the sum of probabilities is smaller than u // In this case, we use the last valid index as the sampled id + if (temp_storage.last_valid_id == -1) { + if (tx == 0) { + output[bx] = 0; + valid[bx] = false; + } + return; + } sampled_id = temp_storage.last_valid_id; } - output[bx] = sampled_id; + if (tx == 0) { + output[bx] = sampled_id; + valid[bx] = true; + } } template __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, float* top_p_arr, - IdType* output, IdType* indices, IdType top_k_val, - float top_p_val, uint32_t d, uint64_t philox_seed, - uint64_t philox_offset) { + IdType* output, bool* valid, IdType* indices, + IdType top_k_val, float top_p_val, uint32_t d, + uint64_t philox_seed, uint64_t philox_offset) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; @@ -1156,6 +1199,7 @@ __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, int sampled_id; do { temp_storage.sampled_id = d; + temp_storage.last_valid_id = -1; __syncthreads(); float u = hiprand_uniform(&state) * q; aggregate = 0; @@ -1179,6 +1223,13 @@ __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, // NOTE(Zihao): this would happen when u is very close to 1 // and the sum of probabilities is smaller than u // In this case, we use the last valid index as the sampled id + if (temp_storage.last_valid_id == -1) { + if (tx == 0) { + output[bx] = 0; + valid[bx] = false; + } + return; + } sampled_id = temp_storage.last_valid_id; } double pivot_0 = probs[row_idx * d + sampled_id]; @@ -1242,6 +1293,7 @@ __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, __syncthreads(); if (tx == 0) { output[bx] = sampled_id; + valid[bx] = true; } } @@ -1359,16 +1411,16 @@ hipError_t SamplingFromLogits(T* logits, IdType* output, IdType* indices, uint32 } template -hipError_t SamplingFromProb(T* probs, IdType* output, IdType* indices, uint32_t batch_size, - uint32_t d, bool deterministic, uint64_t philox_seed, - uint64_t philox_offset, hipStream_t stream = 0) { +hipError_t SamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, + uint32_t batch_size, uint32_t d, bool deterministic, + uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &indices, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &output, &valid, &indices, &d, &philox_seed, &philox_offset}; const uint32_t smem_size = sizeof(SamplingTempStorage); DISPATCH_ALIGNED_VEC_SIZE( @@ -1382,8 +1434,8 @@ hipError_t SamplingFromProb(T* probs, IdType* output, IdType* indices, uint32_t } template -hipError_t TopKSamplingFromProb(T* probs, IdType* output, IdType* indices, T* top_k_arr, - uint32_t batch_size, uint32_t top_k_val, uint32_t d, +hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, + T* top_k_arr, uint32_t batch_size, uint32_t top_k_val, uint32_t d, bool deterministic, uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); @@ -1393,7 +1445,7 @@ hipError_t TopKSamplingFromProb(T* probs, IdType* output, IdType* indices, T* to const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &indices, &top_k_arr, + void* args[] = {&probs, &output, &valid, &indices, &top_k_arr, &top_k_val, &d, &philox_seed, &philox_offset}; DISPATCH_ALIGNED_VEC_SIZE( @@ -1409,9 +1461,9 @@ hipError_t TopKSamplingFromProb(T* probs, IdType* output, IdType* indices, T* to } template -hipError_t TopPSamplingFromProb(T* probs, IdType* output, IdType* indices, T* top_p_arr, - uint32_t batch_size, T top_p_val, uint32_t d, bool deterministic, - uint64_t philox_seed, uint64_t philox_offset, +hipError_t TopPSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, + T* top_p_arr, uint32_t batch_size, T top_p_val, uint32_t d, + bool deterministic, uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); @@ -1420,7 +1472,7 @@ hipError_t TopPSamplingFromProb(T* probs, IdType* output, IdType* indices, T* to const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &indices, &top_p_arr, + void* args[] = {&probs, &output, &valid, &indices, &top_p_arr, &top_p_val, &d, &philox_seed, &philox_offset}; DISPATCH_ALIGNED_VEC_SIZE( @@ -1436,8 +1488,8 @@ hipError_t TopPSamplingFromProb(T* probs, IdType* output, IdType* indices, T* to } template -hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, IdType* indices, - uint32_t batch_size, float min_p_val, uint32_t d, +hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, bool* valid, + IdType* indices, uint32_t batch_size, float min_p_val, uint32_t d, bool deterministic, uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); @@ -1447,7 +1499,7 @@ hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, IdType* const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &min_p_arr, &output, &indices, + void* args[] = {&probs, &min_p_arr, &output, &valid, &indices, &min_p_val, &d, &philox_seed, &philox_offset}; DISPATCH_ALIGNED_VEC_SIZE( @@ -1464,8 +1516,8 @@ hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, IdType* template hipError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, IdType* output, - IdType* indices, uint32_t batch_size, IdType top_k_val, - T top_p_val, uint32_t d, bool deterministic, + bool* valid, IdType* indices, uint32_t batch_size, + IdType top_k_val, T top_p_val, uint32_t d, bool deterministic, uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); @@ -1475,7 +1527,7 @@ hipError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, I const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &top_k_arr, &top_p_arr, &output, &indices, + void* args[] = {&probs, &top_k_arr, &top_p_arr, &output, &valid, &indices, &top_k_val, &top_p_val, &d, &philox_seed, &philox_offset}; DISPATCH_ALIGNED_VEC_SIZE( From ec984f87638a92e91f6f2a62f33a1e2c2cbfbae3 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 14:51:28 -0400 Subject: [PATCH 3/7] rocm/sampling: accept seed/offset tensors, and honour a per-row one sampling.py has always accepted seed as an int or a 1-D tensor of length 1 or batch_size; ROCm rejected every tensor as "per-request seed/offset tensors are not supported". They are now accepted, so nothing CUDA takes raises here. Upstream's kernel reads seed_arr[0] whatever the length, so a length-batch_size tensor of distinct seeds silently seeds every row from element 0 on CUDA. ROCm reads seed_arr[bx * stride], stride 0 for length 1 and 1 for per-row, so a caller who built a per-row tensor gets per-row seeding. Deliberate divergence: sample streams differ from CUDA for that input, and only that input, because it is the case where CUDA drops the caller's data. Pinned by test_a_per_row_seed_tensor_is_honoured_not_collapsed -- forcing seed_arr[0] fails all six entry points. Indexing is by bx, not row_idx: the tensor is validated against the output length, while row_idx indexes probs and may repeat or skip rows under `indices`, so seed_arr[row_idx] could read past the end. The six scalars are bundled into a by-value PhiloxArgs rather than passed individually. The launchers marshal through untyped `void* args[]` that nothing type-checks, so this keeps one slot per kernel instead of adding four -- the plan called for the flat form; the struct is a deviation taken to shrink exactly that hazard. make_philox casts with static_cast rather than data_ptr(), which TORCH_CHECK-fails on the torch.uint64 tensors sampling.py:718 accepts. --- csrc/rocm/sampling.cu | 63 ++++++++++-------- include/flashinfer/rocm/sampling.cuh | 99 ++++++++++++++++------------ tests/rocm/test_sampling.py | 81 +++++++++++++++++++---- 3 files changed, 157 insertions(+), 86 deletions(-) diff --git a/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 581ebddcfa..291f022a19 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -29,13 +29,27 @@ using namespace flashinfer; // added a `valid` output. The ROCm kernels have neither, so both are absorbed // here rather than in the kernels. -// One scalar seed covers the whole batch; a per-request tensor has no kernel to -// route to, so reject it instead of silently sampling every row from seed_val. -inline void reject_per_request_seed(const std::optional& maybe_seed_arr, - const std::optional& maybe_offset_arr) { - TORCH_CHECK(!maybe_seed_arr.has_value() && !maybe_offset_arr.has_value(), - "per-request seed/offset tensors are not supported on ROCm; " - "pass int seed/offset instead"); +// sampling.py validates length in {1, batch_size} and dtype in {int64, uint64}; +// stride 0 broadcasts the length-1 case. data_ptr() would TORCH_CHECK +// on a uint64 tensor, so the cast is unchecked, as upstream's is. +inline sampling::PhiloxArgs make_philox(const std::optional& maybe_seed_arr, + int64_t seed_val, + const std::optional& maybe_offset_arr, + int64_t offset_val) { + sampling::PhiloxArgs philox{}; + philox.seed_val = static_cast(seed_val); + philox.offset_val = static_cast(offset_val); + if (maybe_seed_arr.has_value()) { + CHECK_INPUT(maybe_seed_arr.value()); + philox.seed_arr = static_cast(maybe_seed_arr->data_ptr()); + philox.seed_stride = maybe_seed_arr->size(0) == 1 ? 0u : 1u; + } + if (maybe_offset_arr.has_value()) { + CHECK_INPUT(maybe_offset_arr.value()); + philox.offset_arr = static_cast(maybe_offset_arr->data_ptr()); + philox.offset_stride = maybe_offset_arr->size(0) == 1 ? 0u : 1u; + } + return philox; } // The kernels write `valid` per row now; this only checks the buffer they write @@ -82,15 +96,13 @@ void sampling_from_logits(at::Tensor logits, at::Tensor output, unsigned int batch_size = output.size(0); unsigned int vocab_size = logits.size(1); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::SamplingFromLogits( static_cast(logits.data_ptr()), static_cast(output.data_ptr()), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + batch_size, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "SamplingFromLogits failed with error code " + std::string(hipGetErrorString(status))); } @@ -105,7 +117,6 @@ void sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor valid, unsigned int batch_size = output.size(0); unsigned int vocab_size = probs.size(1); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); @@ -114,8 +125,8 @@ void sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor valid, static_cast(probs.data_ptr()), static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + batch_size, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "SamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -133,7 +144,6 @@ void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_top_p_arr = maybe_top_p_arr.has_value(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); @@ -143,8 +153,8 @@ void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, - top_p_val, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + top_p_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "TopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -165,7 +175,6 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_top_k_arr = maybe_top_k_arr.has_value(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); @@ -175,8 +184,8 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, - top_k_val, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + top_k_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "TopKSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -197,7 +206,6 @@ void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_min_p_arr = maybe_min_p_arr.has_value(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); @@ -207,8 +215,8 @@ void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v has_min_p_arr ? static_cast(maybe_min_p_arr->data_ptr()) : nullptr, static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, min_p_val, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + batch_size, min_p_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "MinPSamplingFromProb failed with error code " + std::string(hipGetErrorString(status))); } @@ -232,7 +240,6 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te bool has_top_k_arr = maybe_top_k_arr.has_value(); bool has_top_p_arr = maybe_top_p_arr.has_value(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); @@ -244,7 +251,7 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, top_k_val, top_p_val, vocab_size, deterministic, - static_cast(philox_seed), static_cast(philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "TopKTopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -274,8 +281,6 @@ void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_i CHECK_EQ(batch_size, output_accepted_token_num.size(0)); CHECK_EQ(batch_size, output_emitted_draft_token_num.size(0)); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::ChainSpeculativeSampling( @@ -283,8 +288,8 @@ void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_i static_cast(target_probs.data_ptr()), static_cast(output_token_ids.data_ptr()), static_cast(output_accepted_token_num.data_ptr()), static_cast(output_emitted_draft_token_num.data_ptr()), batch_size, - num_speculate_tokens, vocab_size, deterministic, static_cast(philox_seed), - static_cast(philox_offset), stream); + num_speculate_tokens, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); TORCH_CHECK(status == hipSuccess, "ChainSpeculativeSampling failed with error code " + std::string(hipGetErrorString(status))); diff --git a/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index 69c6c9abfb..ec665413b4 100644 --- a/include/flashinfer/rocm/sampling.cuh +++ b/include/flashinfer/rocm/sampling.cuh @@ -703,6 +703,30 @@ struct DataAndIndex { } }; +// Seed/offset as upstream passes them -- a scalar, or a device tensor of length +// 1 or batch_size. Bundled so each kernel keeps one launch-argument slot: the +// launchers marshal through untyped `void* args[]`, which nothing type-checks. +// +// Diverges from CUDA deliberately: upstream reads seed_arr[0] whatever the +// length, so a per-row tensor seeds every row identically there. `stride` is 0 +// for a length-1 tensor and 1 for a per-row one, so ROCm honours what the +// caller built. tests/rocm/test_sampling.py pins it. +struct PhiloxArgs { + uint64_t* seed_arr; + uint64_t seed_val; + uint32_t seed_stride; + uint64_t* offset_arr; + uint64_t offset_val; + uint32_t offset_stride; + + __device__ __forceinline__ uint64_t seed(uint32_t bx) const { + return seed_arr ? seed_arr[bx * seed_stride] : seed_val; + } + __device__ __forceinline__ uint64_t offset(uint32_t bx) const { + return offset_arr ? offset_arr[bx * offset_stride] : offset_val; + } +}; + template __device__ __forceinline__ vec_t GenerateGumbelNoise(uint64_t philox_seed, uint64_t philox_offset, @@ -735,7 +759,7 @@ template __global__ void SamplingFromLogitsKernel(DType* logits, IdType* output, IdType* indices, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { + PhiloxArgs philox) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; using SharedMem = typename BlockReduce, BLOCK_THREADS, @@ -752,7 +776,7 @@ __global__ void SamplingFromLogitsKernel(DType* logits, IdType* output, IdType* } vec_t gumbel_noise = GenerateGumbelNoise( - philox_seed, philox_offset, + philox.seed(bx), philox.offset(bx), static_cast(bx * d + (i * BLOCK_THREADS + tx) * VEC_SIZE)); DataAndIndex cur_data[VEC_SIZE]; #pragma unroll @@ -776,10 +800,10 @@ template __global__ void SamplingFromProbKernel(DType* probs, IdType* output, bool* valid, IdType* indices, - uint32_t d, uint64_t philox_seed, uint64_t philox_offset) { + uint32_t d, PhiloxArgs philox) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; - hiprand_init(philox_seed, bx, philox_offset, &state); + hiprand_init(philox.seed(bx), bx, philox.offset(bx), &state); const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; extern __shared__ __align__( @@ -835,12 +859,11 @@ template __global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, bool* valid, IdType* indices, IdType* top_k_arr, uint32_t top_k_val, - uint32_t d, uint64_t philox_seed, - uint64_t philox_offset) { + uint32_t d, PhiloxArgs philox) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; - hiprand_init(philox_seed, bx, philox_offset, &state); + hiprand_init(philox.seed(bx), bx, philox.offset(bx), &state); const uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[bx]; const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; @@ -962,12 +985,11 @@ template __global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, bool* valid, IdType* indices, float* top_p_arr, float top_p_val, - uint32_t d, uint64_t philox_seed, - uint64_t philox_offset) { + uint32_t d, PhiloxArgs philox) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; hiprandStatePhilox4_32_10_t state; - hiprand_init(philox_seed, bx, philox_offset, &state); + hiprand_init(philox.seed(bx), bx, philox.offset(bx), &state); const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; float top_p = (top_p_arr == nullptr) ? top_p_val : top_p_arr[row_idx]; @@ -1083,12 +1105,11 @@ template __global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdType* output, bool* valid, IdType* indices, float min_p_val, - uint32_t d, uint64_t philox_seed, - uint64_t philox_offset) { + uint32_t d, PhiloxArgs philox) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; float p = (min_p_arr == nullptr) ? min_p_val : min_p_arr[bx]; hiprandStatePhilox4_32_10_t state; - hiprand_init(philox_seed, bx, philox_offset, &state); + hiprand_init(philox.seed(bx), bx, philox.offset(bx), &state); const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; extern __shared__ __align__( @@ -1176,11 +1197,11 @@ template hipError_t SamplingFromLogits(T* logits, IdType* output, IdType* indices, uint32_t batch_size, - uint32_t d, bool deterministic, uint64_t philox_seed, - uint64_t philox_offset, hipStream_t stream = 0) { + uint32_t d, bool deterministic, PhiloxArgs philox, + hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&logits, &output, &indices, &d, &philox_seed, &philox_offset}; + void* args[] = {&logits, &output, &indices, &d, &philox}; const uint32_t smem_size = sizeof( typename BlockReduce, BLOCK_THREADS, REDUCE_ALGO>::TempStorage); @@ -1412,15 +1433,15 @@ hipError_t SamplingFromLogits(T* logits, IdType* output, IdType* indices, uint32 template hipError_t SamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, - uint32_t batch_size, uint32_t d, bool deterministic, - uint64_t philox_seed, uint64_t philox_offset, hipStream_t stream = 0) { + uint32_t batch_size, uint32_t d, bool deterministic, PhiloxArgs philox, + hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &valid, &indices, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &output, &valid, &indices, &d, &philox}; const uint32_t smem_size = sizeof(SamplingTempStorage); DISPATCH_ALIGNED_VEC_SIZE( @@ -1436,8 +1457,7 @@ hipError_t SamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indic template hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, T* top_k_arr, uint32_t batch_size, uint32_t top_k_val, uint32_t d, - bool deterministic, uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { + bool deterministic, PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1445,8 +1465,7 @@ hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* i const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &valid, &indices, &top_k_arr, - &top_k_val, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &output, &valid, &indices, &top_k_arr, &top_k_val, &d, &philox}; DISPATCH_ALIGNED_VEC_SIZE( vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { @@ -1463,8 +1482,7 @@ hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* i template hipError_t TopPSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, T* top_p_arr, uint32_t batch_size, T top_p_val, uint32_t d, - bool deterministic, uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { + bool deterministic, PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1472,8 +1490,7 @@ hipError_t TopPSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* i const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &output, &valid, &indices, &top_p_arr, - &top_p_val, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &output, &valid, &indices, &top_p_arr, &top_p_val, &d, &philox}; DISPATCH_ALIGNED_VEC_SIZE( vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { @@ -1490,8 +1507,7 @@ hipError_t TopPSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* i template hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, bool* valid, IdType* indices, uint32_t batch_size, float min_p_val, uint32_t d, - bool deterministic, uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { + bool deterministic, PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1499,8 +1515,7 @@ hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, bool* va const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &min_p_arr, &output, &valid, &indices, - &min_p_val, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &min_p_arr, &output, &valid, &indices, &min_p_val, &d, &philox}; DISPATCH_ALIGNED_VEC_SIZE( vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { @@ -1518,8 +1533,7 @@ template hipError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, IdType* output, bool* valid, IdType* indices, uint32_t batch_size, IdType top_k_val, T top_p_val, uint32_t d, bool deterministic, - uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { + PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1527,8 +1541,8 @@ hipError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, I const uint32_t smem_size = sizeof(SamplingTempStorage); dim3 nblks(batch_size); dim3 nthrs(BLOCK_THREADS); - void* args[] = {&probs, &top_k_arr, &top_p_arr, &output, &valid, &indices, - &top_k_val, &top_p_val, &d, &philox_seed, &philox_offset}; + void* args[] = {&probs, &top_k_arr, &top_p_arr, &output, &valid, + &indices, &top_k_val, &top_p_val, &d, &philox}; DISPATCH_ALIGNED_VEC_SIZE( vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { @@ -2064,11 +2078,11 @@ __global__ void ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token IdType* output_accepted_token_num, IdType* output_emitted_draft_token_num, uint32_t num_speculative_tokens, uint32_t d, - uint64_t philox_seed, uint64_t philox_offset) { + PhiloxArgs philox) { const uint32_t bx = blockIdx.x, tx = threadIdx.x; const uint32_t row_idx = bx; hiprandStatePhilox4_32_10_t curand_state; - hiprand_init(philox_seed, bx, philox_offset, &curand_state); + hiprand_init(philox.seed(bx), bx, philox.offset(bx), &curand_state); extern __shared__ __align__( alignof(SamplingTempStorage)) @@ -2198,8 +2212,7 @@ hipError_t ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token_ids, IdType* output_accepted_token_num, IdType* output_emitted_draft_token_num, uint32_t batch_size, uint32_t num_speculative_tokens, uint32_t d, bool deterministic, - uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { + PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(DType), d); auto compute_capacity = GetCudaComputeCapability(); @@ -2215,8 +2228,8 @@ hipError_t ChainSpeculativeSampling(DType* draft_probs, IdType* draft_token_ids, &output_emitted_draft_token_num, &num_speculative_tokens, &d, - &philox_seed, - &philox_offset}; + &philox}; + DISPATCH_ALIGNED_VEC_SIZE( vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { auto kernel = ChainSpeculativeSampling Date: Wed, 2 Sep 2026 15:36:18 -0400 Subject: [PATCH 4/7] rocm/renorm: run top-k renorm and mask natively at fp16/bf16 The wrapper upcast the whole tensor and copied the result back, on the stated grounds that "ROCm's kernels are float32-only". They are not: all three are DType-templated and read through cast_load. Three things blocked half, none of them the kernel maths: - TopKMaskLogitsKernel and TopKRenormProbKernel stored through vec_t's float-only `store` overload (vec_dtypes.h:829), so half did not compile - GetMaxValue and GetMinMaxValue hardcoded `float* in_data` while already loading via cast_load Converting the two stores to cast_store and templating the two helpers on the pointer type is the whole change; DType is deduced and stays last, so no call site names it. A half call now saves a full-tensor temporary and two casts. top_p_renorm_probs stays fp32 and says so: sampling.py:490 casts before calling it, so a half tensor cannot reach the op through the public API, and upstream has no dispatch there either. Asserted at the op rather than the wrapper, which would hide a half tensor reaching a kernel that no longer upcasts. check_renorm_io replaces as_fp32 and adds the in/out dtype-equality check that as_fp32 used to guarantee by construction -- CHECK_SHAPE compares sizes only. test_every_row_reports_valid is rewritten: it asserted valid.all() on rows that were all valid, which its own docstring admitted would pass with the fill deleted. It now zeroes every other row and checks both directions. --- csrc/rocm/renorm.cu | 80 ++++++++++++++++------------ include/flashinfer/rocm/sampling.cuh | 16 +++--- tests/rocm/test_sampling.py | 50 ++++++++++++++++- 3 files changed, 104 insertions(+), 42 deletions(-) diff --git a/csrc/rocm/renorm.cu b/csrc/rocm/renorm.cu index d76a93c612..7756b5dca3 100644 --- a/csrc/rocm/renorm.cu +++ b/csrc/rocm/renorm.cu @@ -13,37 +13,30 @@ using namespace flashinfer; // ternary-search kernels, which need neither and are deterministic already, // so those parameters are accepted to match the schema and left unused. -// ROCm's kernels are float32-only. v0.6.18 dropped the wrapper's probs.float() -// for the two top_k ops, so a half input reached the kernel and was read at a -// float stride -- a silent 2x overrun of both buffers. Upcast rather than -// reject, which is what the wrapper did up to 0.5.3. top_p still casts in -// Python; it goes through here anyway so the next sync cannot reopen this. -struct Fp32Pair { - at::Tensor in, out; - bool cast_back; -}; - -inline Fp32Pair as_fp32(const at::Tensor& in, const at::Tensor& out) { - // Before the branch, so the fp32 fast path -- which hands the caller's buffer - // straight to the kernel -- is checked as strictly as the cast path. +// The two top-k ops instantiate their kernels at the caller's dtype rather than +// upcasting the whole tensor. They were fp32-only for a different reason than +// the old comment gave: the kernels are DType-templated, but two of them stored +// through vec_t's float-only `store` overload, so half would not compile. +// +// top_p_renorm_probs stays fp32: sampling.py casts before calling it, so a half +// tensor cannot reach here, and upstream has no dispatch there either. +inline void check_renorm_io(const at::Tensor& in, const at::Tensor& out) { CHECK_INPUT(out); CHECK_SHAPE(in, out); + TORCH_CHECK(in.scalar_type() == out.scalar_type(), "input and output dtype must match, got ", + in.scalar_type(), " and ", out.scalar_type()); TORCH_CHECK(in.scalar_type() == at::kFloat || in.scalar_type() == at::kHalf || in.scalar_type() == at::kBFloat16, "expected float32, float16 or bfloat16, got ", in.scalar_type()); - if (in.scalar_type() == at::kFloat && out.scalar_type() == at::kFloat) { - return {in, out, false}; - } - at::Tensor in_f = in.to(at::kFloat); - // Member form: at::empty_like is not visible in this TU under -xhip. - return {in_f, in_f.new_empty(in_f.sizes()), true}; } void top_p_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, std::optional maybe_top_p_arr, double top_p_val, bool is_deterministic, at::Tensor workspace) { CHECK_INPUT(probs); - auto fp32 = as_fp32(probs, renorm_probs); + check_renorm_io(probs, renorm_probs); + TORCH_CHECK(probs.scalar_type() == at::kFloat, + "top_p_renorm_probs is fp32 on ROCm; sampling.py casts before calling it"); auto device = probs.device(); CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) unsigned int batch_size = probs.size(0); @@ -55,19 +48,18 @@ void top_p_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = sampling::TopPRenormProb( - static_cast(fp32.in.data_ptr()), static_cast(fp32.out.data_ptr()), + static_cast(probs.data_ptr()), static_cast(renorm_probs.data_ptr()), has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, top_p_val, vocab_size, stream); TORCH_CHECK(status == hipSuccess, "TopPRenormProb failed with error code " + std::string(hipGetErrorString(status))); - if (fp32.cast_back) renorm_probs.copy_(fp32.out); } void top_k_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, std::optional maybe_top_k_arr, int64_t top_k_val, at::Tensor row_states_buffer) { CHECK_INPUT(probs); - auto fp32 = as_fp32(probs, renorm_probs); + check_renorm_io(probs, renorm_probs); auto device = probs.device(); CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) unsigned int batch_size = probs.size(0); @@ -76,21 +68,31 @@ void top_k_renorm_probs(at::Tensor probs, at::Tensor renorm_probs, const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::TopKRenormProb( - static_cast(fp32.in.data_ptr()), static_cast(fp32.out.data_ptr()), - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, - top_k_val, vocab_size, stream); + hipError_t status = hipSuccess; + if (probs.scalar_type() == at::kFloat) { + status = sampling::TopKRenormProb( + static_cast(probs.data_ptr()), static_cast(renorm_probs.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + top_k_val, vocab_size, stream); + } else { + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(probs.scalar_type(), c_type, [&] { + status = sampling::TopKRenormProb( + static_cast(probs.data_ptr()), static_cast(renorm_probs.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + top_k_val, vocab_size, stream); + return true; + }); + } TORCH_CHECK(status == hipSuccess, "TopKRenormProb failed with error code " + std::string(hipGetErrorString(status))); - if (fp32.cast_back) renorm_probs.copy_(fp32.out); } void top_k_mask_logits(at::Tensor logits, at::Tensor mask_logits, std::optional maybe_top_k_arr, int64_t top_k_val, at::Tensor row_states_buffer) { CHECK_INPUT(logits); - auto fp32 = as_fp32(logits, mask_logits); + check_renorm_io(logits, mask_logits); auto device = logits.device(); CHECK_DIM(2, logits); // logits: (batch_size, vocab_size) unsigned int batch_size = logits.size(0); @@ -99,12 +101,22 @@ void top_k_mask_logits(at::Tensor logits, at::Tensor mask_logits, const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::TopKMaskLogits( - static_cast(fp32.in.data_ptr()), static_cast(fp32.out.data_ptr()), - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, - top_k_val, vocab_size, stream); + hipError_t status = hipSuccess; + if (logits.scalar_type() == at::kFloat) { + status = sampling::TopKMaskLogits( + static_cast(logits.data_ptr()), static_cast(mask_logits.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + top_k_val, vocab_size, stream); + } else { + DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(logits.scalar_type(), c_type, [&] { + status = sampling::TopKMaskLogits( + static_cast(logits.data_ptr()), static_cast(mask_logits.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + top_k_val, vocab_size, stream); + return true; + }); + } TORCH_CHECK(status == hipSuccess, "TopKMaskLogits failed with error code " + std::string(hipGetErrorString(status))); - if (fp32.cast_back) mask_logits.copy_(fp32.out); } diff --git a/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index ec665413b4..e9be48db57 100644 --- a/include/flashinfer/rocm/sampling.cuh +++ b/include/flashinfer/rocm/sampling.cuh @@ -276,8 +276,8 @@ __device__ __forceinline__ void DeterministicInclusiveSum( } template -__device__ __forceinline__ std::tuple GetMinMaxValue(float* in_data, uint32_t row_idx, + typename TempStorage, typename DType> +__device__ __forceinline__ std::tuple GetMinMaxValue(DType* in_data, uint32_t row_idx, uint32_t d, TempStorage& temp_storage) { const uint32_t tx = threadIdx.x; @@ -318,9 +318,11 @@ __device__ __forceinline__ std::tuple GetMinMaxValue(float* in_dat return std::make_tuple(min_val, max_val); } +// DType is deduced from in_data and so stays last: every call site names the +// first four arguments explicitly. The body already reads through cast_load. template -__device__ __forceinline__ float GetMaxValue(float* in_data, uint32_t row_idx, uint32_t d, + typename TempStorage, typename DType> +__device__ __forceinline__ float GetMaxValue(DType* in_data, uint32_t row_idx, uint32_t d, TempStorage& temp_storage) { const uint32_t tx = threadIdx.x; vec_t in_data_vec; @@ -1873,7 +1875,8 @@ __global__ void TopKMaskLogitsKernel(DType* logits, DType* masked_logits, IdType (logits_vec[j] > pivot) ? logits_vec[j] : -cuda::std::numeric_limits::infinity(); } if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { - logits_vec.store(masked_logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + logits_vec.cast_store(masked_logits + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + + tx * VEC_SIZE); } } } @@ -2000,7 +2003,8 @@ __global__ void TopKRenormProbKernel(DType* probs, DType* renormed_prob, IdType* probs_vec[j] = (probs_vec[j] > pivot) ? probs_vec[j] * normalizer : 0; } if ((i * BLOCK_THREADS + tx) * VEC_SIZE < d) { - probs_vec.store(renormed_prob + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + tx * VEC_SIZE); + probs_vec.cast_store(renormed_prob + row_idx * d + i * BLOCK_THREADS * VEC_SIZE + + tx * VEC_SIZE); } } } diff --git a/tests/rocm/test_sampling.py b/tests/rocm/test_sampling.py index f77cf6d997..41b09fb5f6 100644 --- a/tests/rocm/test_sampling.py +++ b/tests/rocm/test_sampling.py @@ -613,8 +613,8 @@ def test_chain_speculative_sampling( "op", [flashinfer.sampling.top_k_renorm_probs, flashinfer.sampling.top_k_mask_logits], ) -def test_half_input_is_upcast_not_read_at_a_float_stride(op, dtype): - """v0.6.18 admits fp16/bf16 here and stopped casting; ROCm's kernels are fp32. +def test_half_input_runs_natively_and_matches_fp32(op, dtype): + """v0.6.18 admits fp16/bf16 here; the kernels now instantiate at that dtype. Unhandled, the kernel walks 4 bytes per element of a 2-byte buffer and writes the overrun back -- silent corruption, not a crash. torch.equal is the @@ -638,6 +638,28 @@ def _seed_case(op): return x, {k: v for k, v in args.items() if k in op} +def test_top_p_renorm_rejects_half_because_python_casts_first(): + """sampling.py casts before calling, so the op is fp32-only by contract. + + Asserted at the op, not the wrapper: the wrapper would hide a half tensor + reaching a kernel that no longer upcasts it. + """ + flashinfer.sampling.get_sampling_module() + probs = torch.rand(4, 128, dtype=torch.float16, device="cuda") + probs /= probs.sum(dim=-1, keepdim=True) + out = torch.empty_like(probs) + + with pytest.raises(RuntimeError, match="fp32 on ROCm"): + torch.ops.sampling.top_p_renorm_probs( + probs, + out, + None, + 0.9, + False, + torch.empty(1, dtype=torch.int32, device="cuda"), + ) + + @pytest.mark.parametrize( "op", [ @@ -741,6 +763,30 @@ def test_a_row_with_no_positive_probability_reports_invalid(op, kwargs): assert not bool(valid.any()), "a row with no positive probability is not valid" +def test_valid_is_written_per_row_not_filled(): + """`valid` used to be fill_(true) before the kernel ran; now the kernel writes it. + + Half the rows have no positive probability, so a fill -- in either direction + -- fails. Pre-filling the opposite of the expected answer is what makes the + write observable. + """ + batch, vocab = 64, 256 + probs = torch.rand(batch, vocab, device="cuda") + probs /= probs.sum(dim=-1, keepdim=True) + probs[1::2] = 0.0 + + flashinfer.sampling.get_sampling_module() + samples = torch.empty(batch, dtype=torch.int32, device="cuda") + valid = torch.zeros(batch, dtype=torch.bool, device="cuda") + torch.ops.sampling.sampling_from_probs( + probs, samples, valid, None, False, None, 0, None, 0 + ) + + assert bool(valid[0::2].all()), "rows that can be sampled must report valid" + assert not bool(valid[1::2].any()), "all-zero rows must report invalid" + assert torch.all((samples >= 0) & (samples < vocab)) + + def test_every_row_reports_valid(): """ROCm's kernels have no reject path, so return_valid is uniformly true. From 487f552936139e39cc184b00cfa42336a24db3a6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 17:47:51 -0400 Subject: [PATCH 5/7] rocm/sampling: validate seed/offset tensors in C++, not only in Python ROCm is the only backend that indexes seed_arr per row, so a tensor shorter than batch_size reads past the end instead of being harmlessly ignored the way upstream's seed_arr[0] does. Leaning on sampling.py's validation was not enough: the raw torch op is reachable without the wrapper, and this PR's own tests call it that way. Checks length in {1, batch_size}, dtype in {int64, uint64}, 1-D, and device against the output tensor -- upstream validates dim/dtype/device in C++ too (csrc/sampling.cu:26); the length check is the ROCm-specific addition. Also corrects the half-precision test's assertion. torch.equal held while the wrapper upcast both sides through the same fp32 kernel; the native path picks vec_size 8 where fp32 picks 4, so the block reduction sums in a different order and the normalizer can differ by an ulp. The input is now seeded too -- it was unseeded, so a straddling value would have failed intermittently in CI rather than reproducibly. The remaining comments that described the old absorb-in-the-wrapper design are updated; two of the three "departures" that section named are gone. --- csrc/rocm/sampling.cu | 55 +++++++++++++++++++++++++------------ tests/rocm/test_sampling.py | 35 +++++++++++++++++++++-- 2 files changed, 70 insertions(+), 20 deletions(-) diff --git a/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 291f022a19..98a9c65a5e 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -25,29 +25,42 @@ typedef hipStream_t cudaStream_t; using namespace flashinfer; -// v0.6.18 replaced the scalar philox pair with optional per-request tensors and -// added a `valid` output. The ROCm kernels have neither, so both are absorbed -// here rather than in the kernels. - -// sampling.py validates length in {1, batch_size} and dtype in {int64, uint64}; +// v0.6.18 replaced the scalar philox pair with optional seed/offset tensors and +// added a `valid` output. Both are handled in the kernels now; this file only +// marshals them. +// // stride 0 broadcasts the length-1 case. data_ptr() would TORCH_CHECK // on a uint64 tensor, so the cast is unchecked, as upstream's is. +inline uint32_t philox_stride(const at::Tensor& t, const char* name, unsigned int batch_size, + const at::Tensor& reference) { + CHECK_INPUT(t); + CHECK_DIM(1, t); + TORCH_CHECK(t.scalar_type() == at::kLong || t.scalar_type() == at::kUInt64, name, + " tensor must be int64 or uint64, got ", t.scalar_type()); + TORCH_CHECK(t.device() == reference.device(), name, " tensor must be on ", reference.device()); + // Checked here, not only in sampling.py: ROCm is the only backend that indexes + // this per row, so a short tensor reads past the end rather than being ignored, + // and the raw torch op is reachable without the wrapper. + TORCH_CHECK(t.size(0) == 1 || t.size(0) == static_cast(batch_size), name, + " tensor length must be 1 or ", batch_size, ", got ", t.size(0)); + return t.size(0) == 1 ? 0u : 1u; +} + inline sampling::PhiloxArgs make_philox(const std::optional& maybe_seed_arr, int64_t seed_val, const std::optional& maybe_offset_arr, - int64_t offset_val) { + int64_t offset_val, unsigned int batch_size, + const at::Tensor& reference) { sampling::PhiloxArgs philox{}; philox.seed_val = static_cast(seed_val); philox.offset_val = static_cast(offset_val); if (maybe_seed_arr.has_value()) { - CHECK_INPUT(maybe_seed_arr.value()); + philox.seed_stride = philox_stride(*maybe_seed_arr, "seed", batch_size, reference); philox.seed_arr = static_cast(maybe_seed_arr->data_ptr()); - philox.seed_stride = maybe_seed_arr->size(0) == 1 ? 0u : 1u; } if (maybe_offset_arr.has_value()) { - CHECK_INPUT(maybe_offset_arr.value()); + philox.offset_stride = philox_stride(*maybe_offset_arr, "offset", batch_size, reference); philox.offset_arr = static_cast(maybe_offset_arr->data_ptr()); - philox.offset_stride = maybe_offset_arr->size(0) == 1 ? 0u : 1u; } return philox; } @@ -102,7 +115,8 @@ void sampling_from_logits(at::Tensor logits, at::Tensor output, static_cast(logits.data_ptr()), static_cast(output.data_ptr()), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "SamplingFromLogits failed with error code " + std::string(hipGetErrorString(status))); } @@ -126,7 +140,8 @@ void sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor valid, valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "SamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -154,7 +169,8 @@ void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, top_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "TopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -185,7 +201,8 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, top_k_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "TopKSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -216,7 +233,8 @@ void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, min_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "MinPSamplingFromProb failed with error code " + std::string(hipGetErrorString(status))); } @@ -251,7 +269,8 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, batch_size, top_k_val, top_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), + stream); TORCH_CHECK(status == hipSuccess, "TopKTopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -289,7 +308,9 @@ void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_i static_cast(output_accepted_token_num.data_ptr()), static_cast(output_emitted_draft_token_num.data_ptr()), batch_size, num_speculate_tokens, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset), stream); + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output_token_ids), + stream); TORCH_CHECK(status == hipSuccess, "ChainSpeculativeSampling failed with error code " + std::string(hipGetErrorString(status))); diff --git a/tests/rocm/test_sampling.py b/tests/rocm/test_sampling.py index 41b09fb5f6..96b5e6c2e3 100644 --- a/tests/rocm/test_sampling.py +++ b/tests/rocm/test_sampling.py @@ -605,7 +605,7 @@ def test_chain_speculative_sampling( assert torch.all(emitted_num + 1 == (output_token_ids != -1).sum(dim=1)) -# --- ROCm's three departures from the v0.6.18 sampling ABI ------------------ +# --- ROCm's remaining divergence from the v0.6.18 sampling ABI -------------- @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -620,12 +620,15 @@ def test_half_input_runs_natively_and_matches_fp32(op, dtype): the overrun back -- silent corruption, not a crash. torch.equal is the assertion that fails without the fix; the dtype only restates the wrapper. """ + torch.manual_seed(0) x = torch.rand(4, 512, device="cuda").to(dtype) got = op(x, 10) assert got.dtype == dtype - # half -> fp32 is lossless, so this is exact, not approximate. - assert torch.equal(got, op(x.float(), 10).to(dtype)) + # Not exact: the native half path picks vec_size 8 where fp32 picks 4, so the + # block reduction sums in a different order and the normalizer can differ by + # an ulp. Same reason the input is seeded -- the straddling case is rare. + torch.testing.assert_close(got.float(), op(x.float(), 10), rtol=1e-2, atol=1e-3) def _seed_case(op): @@ -763,6 +766,32 @@ def test_a_row_with_no_positive_probability_reports_invalid(op, kwargs): assert not bool(valid.any()), "a row with no positive probability is not valid" +def test_a_mismatched_seed_tensor_is_rejected_at_the_op(): + """ROCm indexes seed_arr[bx], so a short tensor would read past the end. + + Asserted at the raw op: sampling.py checks the length too, but the op is + reachable without it, and ROCm is the only backend where the length matters. + """ + flashinfer.sampling.get_sampling_module() + batch, vocab = 8, 128 + probs = torch.rand(batch, vocab, device="cuda") + probs /= probs.sum(dim=-1, keepdim=True) + samples = torch.empty(batch, dtype=torch.int32, device="cuda") + valid = torch.empty(batch, dtype=torch.bool, device="cuda") + + def run(seed): + torch.ops.sampling.sampling_from_probs( + probs, samples, valid, None, False, seed, 0, None, 0 + ) + + with pytest.raises(RuntimeError, match="length must be 1 or 8"): + run(torch.arange(3, dtype=torch.int64, device="cuda")) + with pytest.raises(RuntimeError, match="int64 or uint64"): + run(torch.zeros(batch, dtype=torch.int32, device="cuda")) + + run(torch.arange(batch, dtype=torch.int64, device="cuda")) # the valid shape + + def test_valid_is_written_per_row_not_filled(): """`valid` used to be fill_(true) before the kernel ran; now the kernel writes it. From 517fac6ba08b6ce0076293d6ee9e10b5916518e6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 19:46:51 -0400 Subject: [PATCH 6/7] rocm/sampling: dispatch on the output id dtype instead of assuming int32 sampling.py sizes `samples` as `indices.dtype`, so an int64 `indices` is an ordinary call; upstream handles it with DISPATCH_DLPACK_IDTYPE_TO_CTYPE(output.dtype(), ...) (csrc/sampling.cu:86). ROCm cast to int* regardless, so the kernel wrote two int32 ids into each int64 slot and the caller read them back as one: >>> flashinfer.sampling.sampling_from_probs(probs, indices=idx64) samples dtype: torch.int64 values: [468151435296, 459561500793, 0, 0] # vocab is 128 Silent corruption, not an error. Pre-existing, and found by the review bot on this PR; fixed here rather than deferred because it is the same failure mode the rest of this branch is closing, the kernels are already IdType-templated, and DISPATCH_PYTORCH_IDTYPE_TO_CTYPE already exists. check_id_out additionally rejects an indices/output dtype mismatch and an output on the wrong device -- previously unchecked on the entry points that did not happen to validate it. --- csrc/rocm/sampling.cu | 167 ++++++++++++++++++++++++------------ tests/rocm/test_sampling.py | 32 +++++++ 2 files changed, 144 insertions(+), 55 deletions(-) diff --git a/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 98a9c65a5e..021cd9e198 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -31,6 +31,22 @@ using namespace flashinfer; // // stride 0 broadcasts the length-1 case. data_ptr() would TORCH_CHECK // on a uint64 tensor, so the cast is unchecked, as upstream's is. +// Upstream dispatches these on output.dtype(); ROCm hardcoded int*, so an int64 +// `indices` (sampling.py sizes `samples` as indices.dtype) had two int32s read +// back as one int64 -- silently corrupt token ids, not an error. +inline void check_id_out(const at::Tensor& output, const std::optional& maybe_indices, + const at::Tensor& reference) { + CHECK_INPUT(output); + CHECK_DIM(1, output); + TORCH_CHECK(output.device() == reference.device(), "output must be on ", reference.device()); + if (maybe_indices.has_value()) { + CHECK_INPUT(maybe_indices.value()); + TORCH_CHECK(maybe_indices->scalar_type() == output.scalar_type(), + "indices and output dtype must match, got ", maybe_indices->scalar_type(), " and ", + output.scalar_type()); + } +} + inline uint32_t philox_stride(const at::Tensor& t, const char* name, unsigned int batch_size, const at::Tensor& reference) { CHECK_INPUT(t); @@ -111,12 +127,17 @@ void sampling_from_logits(at::Tensor logits, at::Tensor output, const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::SamplingFromLogits( - static_cast(logits.data_ptr()), static_cast(output.data_ptr()), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::SamplingFromLogits( + static_cast(logits.data_ptr()), static_cast(output.data_ptr()), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + batch_size, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "SamplingFromLogits failed with error code " + std::string(hipGetErrorString(status))); } @@ -131,17 +152,23 @@ void sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor valid, unsigned int batch_size = output.size(0); unsigned int vocab_size = probs.size(1); + check_id_out(output, maybe_indices, probs); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::SamplingFromProb( - static_cast(probs.data_ptr()), static_cast(output.data_ptr()), - valid.data_ptr(), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::SamplingFromProb( + static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + batch_size, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "SamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -159,18 +186,24 @@ void top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_top_p_arr = maybe_top_p_arr.has_value(); + check_id_out(output, maybe_indices, probs); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::TopPSamplingFromProb( - static_cast(probs.data_ptr()), static_cast(output.data_ptr()), - valid.data_ptr(), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, - top_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::TopPSamplingFromProb( + static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, batch_size, + top_p_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "TopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -191,18 +224,24 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_top_k_arr = maybe_top_k_arr.has_value(); + check_id_out(output, maybe_indices, probs); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::TopKSamplingFromProb( - static_cast(probs.data_ptr()), static_cast(output.data_ptr()), - valid.data_ptr(), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, - top_k_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::TopKSamplingFromProb( + static_cast(probs.data_ptr()), static_cast(output.data_ptr()), + valid.data_ptr(), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + top_k_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "TopKSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -223,18 +262,24 @@ void min_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v unsigned int vocab_size = probs.size(1); bool has_min_p_arr = maybe_min_p_arr.has_value(); + check_id_out(output, maybe_indices, probs); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::MinPSamplingFromProb( - static_cast(probs.data_ptr()), - has_min_p_arr ? static_cast(maybe_min_p_arr->data_ptr()) : nullptr, - static_cast(output.data_ptr()), valid.data_ptr(), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, min_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::MinPSamplingFromProb( + static_cast(probs.data_ptr()), + has_min_p_arr ? static_cast(maybe_min_p_arr->data_ptr()) : nullptr, + static_cast(output.data_ptr()), valid.data_ptr(), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + batch_size, min_p_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "MinPSamplingFromProb failed with error code " + std::string(hipGetErrorString(status))); } @@ -258,19 +303,25 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te bool has_top_k_arr = maybe_top_k_arr.has_value(); bool has_top_p_arr = maybe_top_p_arr.has_value(); + check_id_out(output, maybe_indices, probs); check_valid_out(valid, batch_size); const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::TopKTopPSamplingFromProb( - static_cast(probs.data_ptr()), - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, - has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, - static_cast(output.data_ptr()), valid.data_ptr(), - maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - batch_size, top_k_val, top_p_val, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { + status = sampling::TopKTopPSamplingFromProb( + static_cast(probs.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, + has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, + static_cast(output.data_ptr()), valid.data_ptr(), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + batch_size, top_k_val, top_p_val, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "TopKTopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status))); } @@ -302,15 +353,21 @@ void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_i const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); - hipError_t status = sampling::ChainSpeculativeSampling( - static_cast(draft_probs.data_ptr()), static_cast(draft_token_ids.data_ptr()), - static_cast(target_probs.data_ptr()), static_cast(output_token_ids.data_ptr()), - static_cast(output_accepted_token_num.data_ptr()), - static_cast(output_emitted_draft_token_num.data_ptr()), batch_size, - num_speculate_tokens, vocab_size, deterministic, - make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, - output_token_ids), - stream); + hipError_t status = hipSuccess; + DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output_token_ids.scalar_type(), IdType, [&] { + status = sampling::ChainSpeculativeSampling( + static_cast(draft_probs.data_ptr()), + static_cast(draft_token_ids.data_ptr()), + static_cast(target_probs.data_ptr()), + static_cast(output_token_ids.data_ptr()), + static_cast(output_accepted_token_num.data_ptr()), + static_cast(output_emitted_draft_token_num.data_ptr()), batch_size, + num_speculate_tokens, vocab_size, deterministic, + make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, + output_token_ids), + stream); + return true; + }); TORCH_CHECK(status == hipSuccess, "ChainSpeculativeSampling failed with error code " + std::string(hipGetErrorString(status))); diff --git a/tests/rocm/test_sampling.py b/tests/rocm/test_sampling.py index 96b5e6c2e3..d642aba53c 100644 --- a/tests/rocm/test_sampling.py +++ b/tests/rocm/test_sampling.py @@ -766,6 +766,38 @@ def test_a_row_with_no_positive_probability_reports_invalid(op, kwargs): assert not bool(valid.any()), "a row with no positive probability is not valid" +@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize( + "op, kwargs", + [ + ("sampling_from_logits", {}), + ("sampling_from_probs", {}), + ("top_p_sampling_from_probs", {"top_p": 0.9}), + ("top_k_sampling_from_probs", {"top_k": 10}), + ("min_p_sampling_from_probs", {"min_p": 0.1}), + ("top_k_top_p_sampling_from_probs", {"top_k": 10, "top_p": 0.9}), + ], +) +def test_indices_dtype_is_dispatched_not_assumed(op, kwargs, id_dtype): + """sampling.py sizes `samples` as indices.dtype, so int64 is a normal call. + + ROCm cast output/indices to int* regardless, so an int64 buffer read back two + int32s as one id -- e.g. 468151435296 for a 128-wide vocab. Upstream + dispatches on output.dtype(); this asserts ROCm now does too. + """ + torch.manual_seed(0) + vocab = 128 + x = torch.rand(4, vocab, device="cuda") + if "probs" in op: + x /= x.sum(dim=-1, keepdim=True) + indices = torch.arange(4, dtype=id_dtype, device="cuda") + + got = getattr(flashinfer.sampling, op)(x, **kwargs, indices=indices) + + assert got.dtype == id_dtype + assert torch.all((got >= 0) & (got < vocab)), f"out-of-range ids: {got.tolist()}" + + def test_a_mismatched_seed_tensor_is_rejected_at_the_op(): """ROCm indexes seed_arr[bx], so a short tensor would read past the end. From c23488d99062d385276f81f28f0bc97a114203f7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 2 Sep 2026 20:22:06 -0400 Subject: [PATCH 7/7] rocm/sampling: type top_k_arr as int32, and validate the remaining id tensors sampling.py converts every top_k tensor with .int() regardless of the indices dtype, so top_k_arr is always int32. Typing it IdType* in the kernels was wrong in two directions: it read an int32 buffer at an int64 stride once IdType could widen, and one launcher had been casting it to float* long before that. Honest about the evidence: I could not construct an input that distinguishes the two. `k` is narrowed to uint32_t, so an int64 read takes the low half -- which is the first of the two elements it straddles -- and int32/int64 runs agree even on distinct per-row k with the old typing. So this is typed correctness, not a demonstrated miscompute; the test that ships exercises per-row k rather than claiming to pin the dtype. Also closes the validation gaps the same review named: sampling_from_logits had no check_id_out (it has no `valid`, so it missed the guard that rode along with check_valid_out), and chain_speculative_sampling casts four id tensors off output_token_ids' dtype without checking they agree or are on the device. --- csrc/rocm/sampling.cu | 17 ++++++++++++-- include/flashinfer/rocm/sampling.cuh | 11 +++++---- tests/rocm/test_sampling.py | 35 ++++++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 7 deletions(-) diff --git a/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 021cd9e198..c13b2fcd1e 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -125,6 +125,8 @@ void sampling_from_logits(at::Tensor logits, at::Tensor output, unsigned int batch_size = output.size(0); unsigned int vocab_size = logits.size(1); + check_id_out(output, maybe_indices, logits); + const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = hipSuccess; @@ -235,7 +237,7 @@ void top_k_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Tensor v static_cast(probs.data_ptr()), static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, batch_size, top_k_val, vocab_size, deterministic, make_philox(maybe_seed_arr, philox_seed, maybe_offset_arr, philox_offset, batch_size, output), @@ -312,7 +314,7 @@ void top_k_top_p_sampling_from_probs(at::Tensor probs, at::Tensor output, at::Te DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output.scalar_type(), IdType, [&] { status = sampling::TopKTopPSamplingFromProb( static_cast(probs.data_ptr()), - has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, static_cast(output.data_ptr()), valid.data_ptr(), maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, @@ -351,6 +353,17 @@ void chain_speculative_sampling(at::Tensor draft_probs, at::Tensor draft_token_i CHECK_EQ(batch_size, output_accepted_token_num.size(0)); CHECK_EQ(batch_size, output_emitted_draft_token_num.size(0)); + // Four id tensors are cast off output_token_ids' dtype, so they must agree. + CHECK_INPUT(output_token_ids); + for (const at::Tensor& t : + {draft_token_ids, output_accepted_token_num, output_emitted_draft_token_num}) { + CHECK_INPUT(t); + TORCH_CHECK(t.scalar_type() == output_token_ids.scalar_type(), + "all id tensors must share output_token_ids' dtype, got ", t.scalar_type(), " and ", + output_token_ids.scalar_type()); + TORCH_CHECK(t.device() == device, "all id tensors must be on ", device); + } + const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device); auto stream = at::cuda::getCurrentHIPStream(); hipError_t status = hipSuccess; diff --git a/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index e9be48db57..136ee00bcc 100644 --- a/include/flashinfer/rocm/sampling.cuh +++ b/include/flashinfer/rocm/sampling.cuh @@ -860,7 +860,7 @@ template __global__ void TopKSamplingFromProbKernel(DType* probs, IdType* output, bool* valid, - IdType* indices, IdType* top_k_arr, uint32_t top_k_val, + IdType* indices, int32_t* top_k_arr, uint32_t top_k_val, uint32_t d, PhiloxArgs philox) { const uint32_t batch_size = gridDim.x; const uint32_t bx = blockIdx.x, tx = threadIdx.x; @@ -1196,7 +1196,7 @@ __global__ void MinPSamplingFromProbKernel(DType* probs, float* min_p_arr, IdTyp template -__global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, float* top_p_arr, +__global__ void TopKTopPSamplingFromProbKernel(DType* probs, int32_t* top_k_arr, float* top_p_arr, IdType* output, bool* valid, IdType* indices, IdType top_k_val, float top_p_val, uint32_t d, PhiloxArgs philox) { @@ -1458,8 +1458,9 @@ hipError_t SamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indic template hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, - T* top_k_arr, uint32_t batch_size, uint32_t top_k_val, uint32_t d, - bool deterministic, PhiloxArgs philox, hipStream_t stream = 0) { + int32_t* top_k_arr, uint32_t batch_size, uint32_t top_k_val, + uint32_t d, bool deterministic, PhiloxArgs philox, + hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1532,7 +1533,7 @@ hipError_t MinPSamplingFromProb(T* probs, T* min_p_arr, IdType* output, bool* va } template -hipError_t TopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, IdType* output, +hipError_t TopKTopPSamplingFromProb(T* probs, int32_t* top_k_arr, T* top_p_arr, IdType* output, bool* valid, IdType* indices, uint32_t batch_size, IdType top_k_val, T top_p_val, uint32_t d, bool deterministic, PhiloxArgs philox, hipStream_t stream = 0) { diff --git a/tests/rocm/test_sampling.py b/tests/rocm/test_sampling.py index d642aba53c..76253bca65 100644 --- a/tests/rocm/test_sampling.py +++ b/tests/rocm/test_sampling.py @@ -798,6 +798,41 @@ def test_indices_dtype_is_dispatched_not_assumed(op, kwargs, id_dtype): assert torch.all((got >= 0) & (got < vocab)), f"out-of-range ids: {got.tolist()}" +@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize( + "op, key", + [ + ("top_k_sampling_from_probs", "top_k"), + ("top_k_top_p_sampling_from_probs", "top_k"), + ], +) +def test_a_per_row_top_k_is_read_as_int32_whatever_the_id_dtype(op, key, id_dtype): + """top_k_arr is int32 from sampling.py's .int(), independent of indices dtype. + + Typing it IdType* made an int64 call read the int32 buffer at the wrong + stride, so row i got row i/2's k. Asserted against the scalar-k call, which + shares no code path with the per-row one. + """ + torch.manual_seed(0) + vocab = 128 + probs = torch.rand(4, vocab, device="cuda") + probs /= probs.sum(dim=-1, keepdim=True) + indices = torch.arange(4, dtype=id_dtype, device="cuda") + fn = getattr(flashinfer.sampling, op) + extra = {"top_p": 0.9} if "top_p" in op else {} + + # The k values must differ per row. `k` is narrowed to uint32_t, so an int64 + # read of an int32 buffer takes the low half -- which is the *first* of the + # two elements it straddles. A uniform array therefore reads correctly by + # accident, and only distinct values expose the stride. + ks = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="cuda") + per_row = fn(probs, **{key: ks}, **extra, indices=indices, seed=7, offset=0) + + for row, k in enumerate(ks.tolist()): + want = fn(probs, **{key: k}, **extra, indices=indices, seed=7, offset=0) + assert per_row[row] == want[row], f"row {row} was not sampled with k={k}" + + def test_a_mismatched_seed_tensor_is_rejected_at_the_op(): """ROCm indexes seed_arr[bx], so a short tensor would read past the end.