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/csrc/rocm/sampling.cu b/csrc/rocm/sampling.cu index 052748fbf1..c13b2fcd1e 100644 --- a/csrc/rocm/sampling.cu +++ b/csrc/rocm/sampling.cu @@ -25,26 +25,69 @@ 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. - -// 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"); +// 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. +// 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); + 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, 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()) { + philox.seed_stride = philox_stride(*maybe_seed_arr, "seed", batch_size, reference); + philox.seed_arr = static_cast(maybe_seed_arr->data_ptr()); + } + if (maybe_offset_arr.has_value()) { + philox.offset_stride = philox_stride(*maybe_offset_arr, "offset", batch_size, reference); + philox.offset_arr = static_cast(maybe_offset_arr->data_ptr()); + } + return philox; } -// 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, @@ -82,15 +125,21 @@ 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); + check_id_out(output, maybe_indices, logits); 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); + 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))); } @@ -105,16 +154,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); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + 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()), - 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); + 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))); } @@ -132,17 +188,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(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + 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()), - 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); + 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))); } @@ -163,17 +226,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(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + 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()), - 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); + 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))); } @@ -194,18 +264,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(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + 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()), - 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); + 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))); } @@ -229,19 +305,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(); - reject_per_request_seed(maybe_seed_arr, maybe_offset_arr); - mark_all_valid(valid, batch_size); + 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()), - 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); + 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))); } @@ -271,17 +353,34 @@ 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); + // 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 = 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, static_cast(philox_seed), - static_cast(philox_offset), 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/include/flashinfer/rocm/sampling.cuh b/include/flashinfer/rocm/sampling.cuh index 2cebe7c784..136ee00bcc 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; @@ -703,6 +705,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 +761,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 +778,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 @@ -775,11 +801,11 @@ __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, 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__( @@ -789,6 +815,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,21 +841,31 @@ __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, 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; 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]; @@ -848,6 +885,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 +909,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,19 +978,20 @@ __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, 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]; @@ -963,6 +1009,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 +1033,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 +1098,7 @@ __global__ void TopPSamplingFromProbKernel(DType* probs, IdType* output, IdType* __syncthreads(); if (tx == 0) { output[bx] = sampled_id; + valid[bx] = true; } } @@ -1051,12 +1106,12 @@ 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, 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__( @@ -1101,6 +1156,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,22 +1178,32 @@ __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) { +__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) { 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]; const uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[row_idx]; const float p = top_p_arr == nullptr ? top_p_val : top_p_arr[row_idx]; @@ -1156,6 +1222,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 +1246,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 +1316,7 @@ __global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, __syncthreads(); if (tx == 0) { output[bx] = sampled_id; + valid[bx] = true; } } @@ -1336,15 +1411,15 @@ hipError_t OnlineSoftmax(DType* logits, DType* output, uint32_t batch_size, uint 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); @@ -1359,16 +1434,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, 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, &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( @@ -1382,9 +1457,9 @@ 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, - bool deterministic, uint64_t philox_seed, uint64_t philox_offset, +hipError_t TopKSamplingFromProb(T* probs, IdType* output, bool* valid, IdType* indices, + 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); @@ -1393,8 +1468,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, - &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, { @@ -1409,10 +1483,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, - hipStream_t stream = 0) { +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, PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1420,8 +1493,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, - &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, { @@ -1436,10 +1508,9 @@ 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, - bool deterministic, uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { +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, PhiloxArgs philox, hipStream_t stream = 0) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1447,8 +1518,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, - &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, { @@ -1463,11 +1533,10 @@ 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, - uint64_t philox_seed, uint64_t philox_offset, - hipStream_t stream = 0) { +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) { const uint32_t vec_size = std::gcd(16 / sizeof(T), d); auto compute_capacity = GetCudaComputeCapability(); @@ -1475,8 +1544,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, &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, { @@ -1807,7 +1876,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); } } } @@ -1934,7 +2004,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); } } } @@ -2012,11 +2083,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)) @@ -2146,8 +2217,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(); @@ -2163,8 +2233,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 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): + """A 4-row input plus the kwargs `op` needs, seeded for reproducibility.""" + torch.manual_seed(0) + x = torch.rand(4, 128, device="cuda") + if "probs" in op: + x /= x.sum(dim=-1, keepdim=True) + args = {"top_p": 0.9, "top_k": 10, "min_p": 0.1} + 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", + [ + "sampling_from_logits", + "sampling_from_probs", + "top_p_sampling_from_probs", + "top_k_sampling_from_probs", + "min_p_sampling_from_probs", + "top_k_top_p_sampling_from_probs", + ], +) +def test_a_length_one_seed_tensor_matches_the_scalar_seed(op): + """A device-resident seed is upstream's way to avoid a host sync.""" + x, kwargs = _seed_case(op) + fn = getattr(flashinfer.sampling, op) + + scalar = fn(x, **kwargs, seed=7, offset=0) + tensor = fn( + x, + **kwargs, + seed=torch.tensor([7], dtype=torch.int64, device="cuda"), + offset=torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + assert torch.equal(scalar, tensor) @pytest.mark.parametrize( @@ -639,24 +700,188 @@ def test_half_input_is_upcast_not_read_at_a_float_stride(op, dtype): "top_k_top_p_sampling_from_probs", ], ) -def test_a_tensor_seed_is_rejected_rather_than_collapsed(op): - """One scalar seed covers the batch, so a per-request tensor must not be ignored. +def test_a_per_row_seed_tensor_is_honoured_not_collapsed(op): + """ROCm reads seed_arr[bx]; CUDA reads seed_arr[0] whatever the length. - Every entry point, because a collapsed seed is invisible in the output -- the - samples are still valid tokens, just all drawn from one stream. + A deliberate divergence, so this is the test that pins it. Each row must + match the scalar-seeded call for its own seed -- which also rules out the + collapse, since row i would otherwise carry row 0's draw. """ - x = torch.rand(4, 128, device="cuda") + x, kwargs = _seed_case(op) + fn = getattr(flashinfer.sampling, op) + seeds = torch.tensor([11, 22, 33, 44], dtype=torch.int64, device="cuda") + + got = fn(x, **kwargs, seed=seeds, offset=torch.zeros_like(seeds)) + for row, seed in enumerate(seeds.tolist()): + want = fn(x, **kwargs, seed=seed, offset=0) + assert got[row] == want[row], f"row {row} was not seeded from {seed}" + + +@pytest.mark.parametrize( + "op", + ["sampling_from_probs", "top_k_sampling_from_probs"], +) +def test_a_uniform_per_row_seed_tensor_matches_the_scalar(op): + """Length batch_size but every entry equal: [0] and [bx] agree here. + + Separates "honours the stride" from "reads the tensor at all". + """ + x, kwargs = _seed_case(op) + fn = getattr(flashinfer.sampling, op) + seeds = torch.full((4,), 7, dtype=torch.int64, device="cuda") + + assert torch.equal( + fn(x, **kwargs, seed=seeds, offset=torch.zeros_like(seeds)), + fn(x, **kwargs, seed=7, offset=0), + ) + + +@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" + + +@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) - seed = torch.arange(4, dtype=torch.int64, device="cuda") - args = {"top_p": 0.9, "top_k": 10, "min_p": 0.1} - kwargs = {k: v for k, v in args.items() if k in op} + 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()}" + + +@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. - with pytest.raises(RuntimeError, match="not supported on ROCm"): - getattr(flashinfer.sampling, op)( - x, **kwargs, seed=seed, offset=torch.zeros_like(seed) + 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. + + 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. + + 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.