mudler / mudler/vllm.cpp

cuda_sample.cu: the two remaining grow-only argmax scratches free under queued kernels

Open
#2,916 0 comments 0 reactions 0 assignees View on GitHub

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) launches ArgmaxPartialKernel then
    ArgmaxFinalKernel against g_argmax_scratch.val / .idx and returns without
    synchronizing.
  • RandomSampleCuda (:354-357) does the same against g_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;
  • cudaFree is 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.