deepseek-ai / deepseek-ai/DeepSelect

Top-K vector accesses exceed accepted tensor boundaries

Open
#5 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Cuda
Stars
343
Forks
21
Avg merge
10m
Merged PRs (30d)
2

Description

The stride checks accept output views whose logical rows end partway through a vector store. For example, an int32 output with shape `(2, 1)` and stride `(8, 1)` can have only 36 bytes of backing storage. The shortcut epilogue writes 32 bytes at each row start on SM100, so the second store ends at byte 64. With larger backing storage, the same stores overwrite elements outside the supplied output view.

The input descriptor has a related allocation-boundary issue: it rounds each row to 128 bytes, but the API checks only row stride. A float32 input with shape `(2, 257)`, stride `(512, 1)`, and 3,076 bytes of storage passes the stride check, while the descriptor exposes bytes through offset 3,199. TMA boundary fill cannot protect bytes that the descriptor declares in bounds.

These are source-derived counterexamples at revision `8e70df7`. The relevant code is `EpilogueRunner::topk_select_epilogue`, `TopkSelectKernelBase::make_topk_tensor_map`, and the stride validation in `csrc/api.cpp`.

The output API should write only the requested elements, including for small or non-vector-width `topk`. The input boundary should either provide safe tail loads or reject storage that does not include the descriptor's final padded row. Pointer alignment also needs checking independently of row stride, since a sliced tensor can preserve its stride while shifting its data pointer.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the supplied output and input counterexamples at revision 8e70df7, then inspect EpilogueRunner::topk_select_epilogue, TopkSelectKernelBase::make_topk_tensor_map, and the stride validation in csrc/api.cpp. Trace the requested view bounds, padded descriptor bounds, and pointer alignment. Done means invalid storage is rejected or safely accessed, and valid small or non-vector-width topk writes stay within the supplied output view.

Written by the indexing model from the issue text.

Assessment

Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.