Fix uninitialised last_valid_id, honour seed tensors, run half renorm natively - #354
Conversation
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.
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.
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<int64_t>(), which TORCH_CHECK-fails on the torch.uint64 tensors sampling.py:718 accepts.
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.
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.
There was a problem hiding this comment.
🟡 Changes recommended
ROCm sampling C++ entry points still assume int32 output/indices (despite Python wrappers allowing int64), which can silently corrupt token IDs, and some paths still don’t validate output device/shape consistency.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR addresses ROCm sampling ABI divergences introduced around the v0.6.18 sync by fixing an uninitialized shared-memory sentinel (last_valid_id), adding proper support for seed/offset tensors (including per-row seeding), and enabling native fp16/bf16 execution for the top-k renorm/mask kernels.
Changes:
- Initialize and guard
last_valid_idin ROCm sampling kernels, and make kernels writevalidper-row. - Add
PhiloxArgsplumbing so ROCm honors scalar vs tensor seeds/offsets (including per-row stride semantics). - Remove forced fp32 upcasting for top-k renorm/mask on ROCm by fixing dtype-templated stores/loads; add/adjust ROCm tests to pin behavior.
File summaries
| File | Description |
|---|---|
| tests/rocm/test_sampling.py | Adds regression + behavior tests for invalid rows, seed tensor semantics, and native half/bf16 renorm behavior. |
| include/flashinfer/rocm/sampling.cuh | Introduces PhiloxArgs, initializes last_valid_id, writes valid in-kernel, and enables half/bf16 stores via cast_store. |
| csrc/rocm/sampling.cu | Marshals seed/offset tensors into PhiloxArgs, wires valid outputs into kernel launches, and updates related argument lists. |
| csrc/rocm/renorm.cu | Keeps top-p renorm fp32-only explicitly; dispatches top-k renorm/mask to run natively in fp16/bf16 without whole-tensor upcasts. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
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.
There was a problem hiding this comment.
🟡 Changes recommended
The ROCm C++ sampling wrappers still have validated-by-wrapper gaps (notably for chain speculative outputs and top-k-per-row dtype contract), which can lead to incorrect pointer interpretation or unsafe kernel launches via raw torch.ops calls.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
csrc/rocm/sampling.cu:357
- chain_speculative_sampling dispatches IdType from output_token_ids.scalar_type() and then casts draft_token_ids/output_accepted_token_num/output_emitted_draft_token_num to IdType* without validating those tensors’ dtypes or even CHECK_INPUT()ing the output tensors (CUDA/contiguity). A raw torch.ops call with mismatched dtypes or CPU outputs will lead to the kernel reading/writing through the wrong pointer type/address.
const at::cuda::OptionalHIPGuardMasqueradingAsCUDA device_guard(device);
auto stream = at::cuda::getCurrentHIPStream();
hipError_t status = hipSuccess;
DISPATCH_PYTORCH_IDTYPE_TO_CTYPE(output_token_ids.scalar_type(), IdType, [&] {
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
… 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.
|
Suppressed comment from review 5096415407 ( One correction on the inline
|
There was a problem hiding this comment.
🟡 Changes recommended
The ROCm sampling kernels still index per-request top_p/top_k tensors using row_idx when indices is provided, and the C++ wrapper lacks some critical shape/device validation for raw-op callers, both of which can lead to incorrect behavior or out-of-bounds access.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
include/flashinfer/rocm/sampling.cuh:1209
- TopKTopPSamplingFromProbKernel indexes top_k_arr/top_p_arr using row_idx (indices[bx]) rather than bx. This makes top_k/top_p effectively per-distribution instead of per-request (contradicting the Python wrapper, which builds these tensors with length=batch_size=indices.size(0)) and can go out of bounds if indices refer to a probs row >= batch_size.
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];
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
|
Suppressed comment from review 5096576919 ( Closing the loop. Three rounds, and rounds 1 and 2 were worth it — round 1 found a genuine pre-existing bug (int64 Where I have drawn the line: Later rounds re-raising the |
Summary
Closes the three sampling divergences the v0.6.18 sync left behind. Reading the code rather than the notes changed all three: one was a live memory bug, and the other two were blocked by less than their comments claimed.
What changed
1.
last_valid_idwas never initialised — a wild read, not a missing featurecsrc/rocm/sampling.cusaid the ROCm kernels "have no reject path", sovalidwas filledtruebefore launch. They also had no sentinel:last_valid_idwas written in exactly one place, conditionally (if (tx == 0 && max_valid_index != -1)), and initialised nowhere. Upstream initialises it at five sites. Living in dynamicextern __shared__, it returned whatever the previous block left in LDS whenever no element qualified.A row with no positive probability reproduces it directly:
top_k_sampling_from_probsnever got that far — it fed the garbage intoprobs[row_idx * d + sampled_id]and aborted the process.This is why the fix could not be a guard on
-1: there was no-1to guard on. The initialisations are the load-bearing half; the guards then mirror upstream (output=0,valid=false).2. Seed tensors were rejected, and upstream's own handling is not what it looks like
sampling.pyhas always acceptedseedas an int or a 1-D tensor of length 1 orbatch_size; ROCm raised on every tensor. Upstream's kernel readsseed_arr[0]unconditionally, so on CUDA a length-batch_sizetensor of distinct seeds silently seeds every row from element 0.ROCm now accepts everything CUDA accepts, and reads
seed_arr[bx * stride]— stride 0 for length 1, 1 for per-row. Deliberate divergence: a caller who builds a per-row tensor gets per-row seeding here and does not on CUDA, so sample streams differ for that input and only that input — the case where CUDA drops the caller's data.3. The half renorm kernels were three small things away from compiling
renorm.cuupcast the whole tensor because "ROCm's kernels are float32-only". All three areDType-templated and read throughcast_load. What actually blocked half:TopKMaskLogitsKernelandTopKRenormProbKernelstored throughvec_t's float-onlystoreoverloadGetMaxValueandGetMinMaxValuehardcodedfloat* in_datawhile already loading viacast_loadConverting two stores to
cast_storeand templating two helpers on the pointer type is the whole change.top_p_renorm_probsstays fp32 and now says so —sampling.py:490casts before calling it, so a half tensor cannot reach it, and upstream has no dispatch there either.Notes on how this was built
The launchers marshal kernel arguments through untyped
void* args[]arrays that nothing type-checks against the kernel signature, so a mis-ordered slot is silent memory corruption rather than a build error. Each kernel, its launcher, its args array and its C++ call site were changed and built one at a time. The six philox values are bundled into a by-valuePhiloxArgsso each kernel keeps one launch slot rather than gaining four — a deviation from the approved plan, taken to shrink exactly that hazard.make_philoxvalidates length, dtype, rank and device in C++, not only insampling.py: ROCm is the only backend that indexes the seed per row, so a short tensor reads past the end, and the raw torch op is reachable without the wrapper.chain_speculative_samplingis untouched — upstream has novalidoutput there and never initialises the sentinel either, so both platforms read uninitialised LDS. Inherited deliberately; there is nowhere in that API to report it.Test plan
valid, never a sentinel, and uses a 256-row batch so LDS is recycled across blocks — a single block can read a plausible value and hide the bugseed_arr[0](CUDA's behaviour) fails all six entry points oftest_a_per_row_seed_tensor_is_honoured_not_collapsedtests/rocm, 29254 passedpre-commit run -aFour failures in the gfx942 run (
test_amd_coverage.py::TestManifest,test_redistribution_licences.py) are environmental, not regressions: they fail identically on the base at3947c0bb4becausegit -C /wt ls-filesexits 128 when a linked worktree's.gitfile points outside the bind mount.Verified at
be130758a, then rebased onto3947c0bb4; that base moved onlytests/rocm/test_git_describe.py, which this change does not touch.Not in this change
min_p_sampling_from_probsis guarded like the others but is not reachable via an all-zero row — its predicate isx >= pivotwithpivot = max_val * p, so every element qualifies. Its guard is exercised by the A/B rather than by an input.