Skip to content

feat: build a DeepSelect wheel (sm_90a + sm_100a + sm_103a) - #17

Open
faresobeid wants to merge 1 commit into
mainfrom
feat/deep-select
Open

faresobeid wants to merge 1 commit into
mainfrom
feat/deep-select

Conversation

@faresobeid

@faresobeid faresobeid commented Oct 10, 2026 •

Copy link
Copy Markdown

Adds a deep-select wheel to the kernel release, built from deepseek-ai/DeepSelect @ bfa4507 (v1.0.0 + Ascend).

prime-rl wants it for the GLM-5.x / DeepSeek-V4 DSA indexer top-k (fp8_indexer.py, today torch.topk over fp32 rows up to 131k wide).

  • Upstream compiles only sm_100a / sm_103a. The sm_90a gencode line is commented out upstream only to save compile time. The script turns it on, so one wheel covers H100/H200 and B200/B300. There is no sm_120.
  • setup.py calls lspci before it reads DEEP_SELECT_BUILD_TARGET_PLATFORM. Build containers have neither lspci nor a GPU, so the script puts a one-line lspci stub on PATH.
  • Smoke test: deep_select.topk vs torch.topk (same index sets) on sm_90/100/103, skipped elsewhere.

Validated on H200 (CUDA 13.0, torch 2.13.0+cu130):

  • The script builds the wheel in ~5 min with MAX_JOBS=64 (no register spills). cuobjdump shows sm_90a, sm_100a and sm_103a cubins.
  • scripts/smoke_test.py deep-select passes.
  • At GLM-5.3 shapes ([32k x 131k] fp32, top-2048): torch.topk 54 ms, DeepSelect 11-12 ms (3.6-10.9 ms with per-row end). The selected values match torch.topk exactly in every row; index sets differ only where logits tie at the k boundary.
  • Not validated: the aarch64 build, and running on SM100/SM103 (no Blackwell here).

Consumer PR in prime-rl: PrimeIntellect-ai/prime-rl#3958 (it needs a release of this wheel before it can pin it).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant