ROCm split-phase random sample: phase B launches 1024 threads onto 128-entry shared arrays
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
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 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