[Bugfix] Fix candidate buffer overflow in deepseek_v32 topk selector - #3280
liuyun7345 wants to merge 1 commit into
Conversation
tl_topk_impl staged radix candidates in a fixed [2, 4096] shared buffer, but nothing bounds the threshold bucket selected by the first radix pass to 4096 elements. Inputs where more candidates share that bucket (e.g. FP32 scores with common low mantissa bits, or all-equal inputs) made the atomic appends write past the buffer end and crash with CUDA_ERROR_ILLEGAL_ADDRESS. Allocate a (batch, 2, seq_len) global-memory workspace in the tl_topk wrapper instead, bounding both radix buffers by seq_len. The public tl_topk() signature is unchanged. Fixes tile-ai#1351 Signed-off-by: liuyun7345 <liuyun7345@sina.com>
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: tile-ai/tilelang/.coderabbit.yaml Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe top-k selector replaces its fixed shared-memory candidate buffer with a workspace sized to the input sequence. A CUDA-gated regression test checks selection results across value patterns, sequence lengths, top-k sizes, and full or partial ranges. ChangesTop-k selection
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix · Severity of issue fixed: Medium Suggested reviewers: Merge Risk: ⚪ Minimal · up to The workspace repair and regression coverage show no established merge-blocking issue. Merge after normal checks. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This branch has not been deployed
1 participant
Summary
tl_topkinexamples/deepseek_v32/topk_selector.pystages radix candidates in a fixed shared buffer:Nothing bounds the threshold bucket chosen by the first radix pass to 4096 elements. Whenever more candidates land there — e.g. FP32 scores that share their low 10 mantissa bits (#1351), or simply
torch.ones(...)with a small enough topk — the atomic appends write past the end of the buffer and the kernel crashes withCUDA_ERROR_ILLEGAL_ADDRESS.Changes
tl_topk_implnow takes aworkspace: T.Tensor[(batch, 2, seq_len), T.int32]argument. Alls_input_idx[...]accesses becomeworkspace[bx, ...], and the fixedSMEM_INPUT_SIZEshared allocation is removed, so both radix buffers are bounded byseq_len.tl_topkwrapper allocates the workspace withtorch.empty((batch, 2, seq_len), dtype=torch.int32, ...)and forwards it. The publictl_topk()signature is unchanged.Validation
RTX 4090D, CUDA 12.1, torch 2.6.0+cu124, tilelang 0.1.14 (wheel):
testing/python/issue/test_tilelang_issue_1351.py— 13 parameterized cases: low-mantissa-bit inputs (positive and negative), all-equal inputs at seq_len 4095/4096/4097/32768, boundary tie buckets (exactly topk-1 larger values followed by an oversized tied bucket), partial[starts, ends)ranges, and random inputs. Selected values are compared against a maskedtorch.topkreference withrtol=0, atol=0(exact FP32 match, order-independent). 13/13 passed.compute-sanitizer --tool memcheck: on the original kernel the first regression case (low-10-mantissa-bit input) already trips the sanitizer —ERROR SUMMARY: 129 errors— and the launch fails withCUDA_ERROR_LAUNCH_FAILED(the run stops there); the patched kernel passes the full 13-case matrix under the sanitizer withERROR SUMMARY: 0 errors.+7-10% from global-memory candidate atomics; the wrapper additionally pays one
torch.empty(~8 µs) for the workspace (batch × 2 × seq_len × 4 B, e.g. 16 MB at batch=64, seq_len=32768).ruff check/ruff format --check/codespellclean on both files (ruff 0.16.8, one minor version above the 0.16.6 CI pin).Fixes #1351
Summary
tl_topk_impl’s fixed 4,096-element candidate buffer with a workspace sized toseq_len.tl_topkallocates and passes the workspace; its public signature remains unchanged._RADIXthreads for ROCm. Other platforms retain 1,024 threads.torch.topkreference.Testing
The changes add the regression test. No test run results are available in the supplied evidence.