cuda_sample.cu: the two remaining grow-only argmax scratches free under queued kernels
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 423
- Forks
- 53
- Avg merge
- 20h 26m
- Merged PRs (30d)
- 310
Description
Row: -
src/vt/cuda/cuda_sample.cu still holds two file-scope grow-only argmax
scratches, and both are the exact shape SPEC-DFLASH2 A2-2 deleted from the
rejection-sampling path.
ArgScratch g_argmax_scratch; // cuda_sample.cu:187
ArgScratch g_sample_scratch; // cuda_sample.cu:188
EnsureArgScratch (:174, called at :230 and :350) grows a scratch by
cudaFreeing the current buffer and cudaMallocing a larger one. Both readers
return with their kernels still queued:
GreedyArgmaxCuda(:233-236) launchesArgmaxPartialKernelthen
ArgmaxFinalKernelagainstg_argmax_scratch.val/.idxand returns without
synchronizing.RandomSampleCuda(:354-357) does the same againstg_sample_scratch.
So a second call that needs more rows frees the buffer the first call's still
queued ArgmaxFinalKernel reads. That is bit-for-bit the argument A2-2 made
about g_reject_argmax:
- the free lands under a queued kernel, so the read is a device use-after-free;
- the loud outcome is an illegal memory access; the quiet one is that the read
lands in another allocation and the wrong token id is emitted, with nothing
raised; cudaFreeis also a device-wide synchronize inside a function whose contract
is that it waits on nothing, so the grow path silently drains the stream.
Unlike the rejection-sampler case this one is on the ORDINARY decode path, not on
an unreached speculative arm, and it is reachable today with two runners in one
process or with any sequence of sample() calls whose row counts increase while
an earlier launch is in flight.
Not a regression, and not caused by A2-2. A2-2 removed the rejection-path
instance (g_reject_argmax) by making the scratch a parameter of
vt::GreedyRejectionSample, owned per call by
vllm::v1::RejectionSamplerDeviceOutput. It did not touch these two, which are
out of that wave's scope. This issue exists so the sibling is named rather than
left for someone to rediscover: the repair deleted one instance of a pattern and
left two identical ones in the same file.
Suggested repair, mirroring what A2-2 did: make the scratch a caller-owned
buffer, or give it a lifetime object whose destructor drains the queue its work
is on before it frees, as
include/vllm/v1/spec_decode/rejection_sampler.h::RejectionSamplerDeviceOutput
now does. A same-size reuse must stay allocation-free; only a GROW is the
dangerous edge.
Evidence needed to close it: a red-before test on a GPU (this class is
invisible on the CPU tier and on a unified-memory backend, where the kernels have
already finished), plus the nvcc build. Neither is available on the CPU box the
A2-2 repair ran on.
Found while repairing review findings on SPEC-DFLASH2 A2-2 (#2802). Listed under
## Owed in .agents/specs/dflash2-async-spec-sampler.md.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in src/vt/cuda/cuda_sample.cu at EnsureArgScratch, GreedyArgmaxCuda, and RandomSampleCuda, then compare the lifetime handling in include/vllm/v1/spec_decode/rejection_sampler.h::RejectionSamplerDeviceOutput. Reproduce the failure with a red-before GPU test, verify same-size reuse stays allocation-free and queued kernels do not use freed storage, then run the nvcc build.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100