mudler / mudler/vllm.cpp

ROCm split-phase random sample: phase B launches 1024 threads onto 128-entry shared arrays

Open
#3,022 1 comment 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: BACKEND-ROCM

RandomSampleSplitBK in src/vt/rocm/rocm_sample.hip declares its reduction
scratch at the partial count:

__shared__ float sh_score[kSampleSplitBlocks];    // 128
__shared__ int64_t sh_idx[kSampleSplitBlocks];    // 128

and is launched at the vocab block width:

RandomSampleSplitBK<<<static_cast<unsigned>(n), kVocabBlock, 0, s>>>(   // 1024

Every thread then executes the stores unconditionally:

sh_score[threadIdx.x] = best_v;
sh_idx[threadIdx.x] = best_j;

The if (threadIdx.x < blocks_per_row) above them guards only the global read
of the partials, not the shared write. Threads 128-1023 write past both arrays.
sh_score is 128 floats, so sh_score[128] lands on sh_idx[0] -- the element
the kernel writes to out[row]. The overflowing threads carry kNegInf and
kArgSentinel, so the store races the real sh_idx[0] and can leave a float
bit pattern where a token id belongs.

The path is default-ON (VT_SAMPLE_SPLIT defaults to enabled) and is selected
for any v >= 4096 && n <= 64, which is every production vocab at low
concurrency. The consequence is a wrong sampled token, not a crash.

Why the gate did not catch it

The split path has no test. ROCm random_sample agrees with CPU on the vast majority of rows in tests/vt/test_ops_sample.cpp runs V = 128, below the
v >= 4096 bar, so it takes the single-block kernel. The distribution case runs
V = 4. #3010 reports test_ops_sample 29/29 and the full suite 225/225 green,
and both numbers are true and neither touches the kernel the change adds.

Fix

Launch phase B with kSampleSplitBlocks threads, which is the width its shared
arrays are sized for and the width its reduction tree assumes, plus a test at
V = 8192, N = 8 that selects the split path.

Found while reviewing #3010. Fixed in the same flow.

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 with RandomSampleSplitBK and its phase-B launch in src/vt/rocm/rocm_sample.hip, then inspect the sampling coverage in tests/vt/test_ops_sample.cpp. Run or extend the sampling test with V = 8192 and N = 8 so it selects the split path. Done means phase B uses the shared-array width and the split-path test validates the sampled result against the CPU behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
backend, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.