[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"project-96162":3},{"id":4,"name":5,"fullName":6,"owner":7,"repo":5,"description":8,"homepage":9,"htmlUrl":9,"language":10,"languages":9,"totalLinesOfCode":9,"stars":11,"forks":12,"watchers":13,"openIssues":14,"contributorsCount":15,"subscribersCount":15,"size":15,"stars1d":15,"stars7d":15,"stars30d":16,"stars90d":15,"forks30d":15,"starsTrendScore":15,"compositeScore":17,"rankGlobal":9,"rankLanguage":9,"license":18,"archived":19,"fork":19,"defaultBranch":20,"hasWiki":21,"hasPages":19,"topics":22,"createdAt":9,"pushedAt":9,"updatedAt":23,"readmeContent":24,"aiSummary":25,"trendingCount":15,"starSnapshotCount":15,"syncStatus":26,"lastSyncTime":27,"discoverSource":28},96162,"DeepSelect","deepseek-ai\u002FDeepSelect","deepseek-ai","DeepSelect: TopK kernels for DeepSeek Sparse Attention (DSA) and Samplers",null,"Cuda",295,14,125,4,0,78,51.33,"MIT License",false,"main",true,[],"2026-09-20 04:01:32","# DeepSelect\n\nDeepSelect is a high performance implementation of the TopK kernel used in DeepSeek Sparse Attention (DSA) (which is used in DeepSeek V3.2, DeepSeek V4, and DeepSeek V4.1 models) and the sampler. It achieves 2 ~ 20x speedup compared to vanilla `torch.topk`.\n\n## News\n\n- 2026.09.10: We've released a brief analysis of the algorithm and its implementation: [English](docs\u002FDeepSelect-deep-dive.md) | [中文](docs\u002FDeepSelect-deep-dive.zh.md)\n- 2026.09.10: We've released DeepSelect v1.0.0\n\n## Supported Cases\n\nTopK workloads vary widely, and the fastest algorithm & implementation highly depends on the input dtype, `batch_size`, `vocab_size`, and `topk`. This repository only focuses on the following cases:\n\n### Lightning Indexer Scenario\n\nThis scenario covers:\n- Input dtype: `torch.bfloat16`\n- `batch_size`: $1 \\sim +\\infty$ (both large and small batch sizes are optimized)\n- `vocab_size`: $1 \\sim +\\infty$ (both large and small vocabularies are optimized)\n- `topk`: small (must be $\\le 4096$; larger values are not supported)\n\nRecommendations:\n- Disable `sorted_index` unless the output has to be ordered by index or by value; enabling either one costs performance.\n- Set `return_value=False` when the values are not needed. This skips the value output and is faster.\n\n### Sampling Scenario\n\nThis scenario covers:\n- Input dtype: `torch.float32`\n- `batch_size`: $1 \\sim +\\infty$\n- `vocab_size`: around 128K\n- `topk`: small (must be $\\le 4096$; larger values are not supported)\n\n## Performance\n\nMeasured with the benchmark in [`tests\u002Ftest.py`](tests\u002Ftest.py)\n(`python3 tests\u002Ftest.py --perf-only`), which reports the ratio against `torch.topk`\non the same input. The metric is effective memory bandwidth: TopK does no\nfloating-point math, so a FLOP rate would not be meaningful here.\n\n### Lightning Indexer Scenario\n\nbfloat16, `topk = 512`, one subplot per batch size, on a shared 0 - 7 TB\u002Fs axis.\n\n![DeepSelect vs torch.topk, bfloat16 Lightning Indexer](assets\u002Fperf_bf16.png)\n\n### Sampling Scenario\n\nfloat32, `vocab_size = 129280`, `topk = 512`.\n\n![DeepSelect vs torch.topk, float32 Sampling](assets\u002Fperf_fp32.png)\n\n## Installation\n\n```bash\ngit clone https:\u002F\u002Fgithub.com\u002Fdeepseek-ai\u002FDeepSelect.git\ncd DeepSelect\ngit submodule update --init --recursive\npip install -v .\n```\n\n## Usage\n\n```python\nimport torch\nimport deep_select\n\n# input: (batch_size, vocab_size), torch.bfloat16 or torch.float32.\n# Its row stride must be a multiple of `deep_select.get_stride_requirement()[0]` bytes, and its last dimension must be contiguous.\nbatch_size, vocab_size, topk = 4, 204800, 1024\n\nx = torch.randn(batch_size, vocab_size, dtype=torch.bfloat16, device=\"cuda\")\n\nvalues, indices = deep_select.topk(\n    x,\n    topk,\n    sorted_index=True,         # return each row's indices in ascending order\n    indices_type=torch.int32,  # torch.int32 or torch.int64\n    return_value=True,         # False skips the value output (~10% faster)\n)\n# values:  (batch_size, topk) of x.dtype\n# indices: (batch_size, topk) of indices_type\n```\n\nThe row stride of the input tensor (`x`) must be aligned to `deep_select.get_stride_requirement()[0]` bytes. For unaligned inputs, padding is necessary.\n\nBoth outputs are allocated by the call, and their strides are aligned to `deep_select.get_stride_requirement()[1]` bytes (so they may be non-contiguous). Pass `output_idx=` to write indices into a buffer you own, and that buffer must satisfy the same stride requirement.\n\nFor the full signature, see [`deep_select\u002Finterface.py`](deep_select\u002Finterface.py).\n\n### Variable-length rows\n\n`end` sets a per-row upper bound (exclusive). Rows shorter than `topk` are padded with\n`value_oob_fill_value` \u002F `idx_oob_fill_value`:\n\n```python\nbatch_size, vocab_size = 2, 129280   # 129280 is a multiple of 256, so float32 is fine\nx = torch.randn(batch_size, vocab_size, dtype=torch.float32, device=\"cuda\")\n\nend = torch.tensor([129280, 100000], dtype=torch.int32, device=\"cuda\")  # (batch_size,)\nvalues, indices = deep_select.topk(x, 1000, end=end, sorted=True,\n                                   indices_type=torch.int64)\n```\n\n### NaN handling\n\nNaN checking is always on. With the default `abort_when_nan_found=True` the kernel invokes `trap()` and aborts. Rows whose length is `\u003C= topk` are never NaN-checked.\n\n## Citation\n\n```text\n@misc{deepselect2026,\n    title={DeepSelect: High-Performance TopK Kernels for DeepSeek Sparse Attention and Sampling},\n    author={Yi Qian and Shengyu Liu and Yichen Li},\n    year={2026},\n    publisher = {GitHub},\n    howpublished = {\\url{https:\u002F\u002Fgithub.com\u002Fdeepseek-ai\u002FDeepSelect}},\n}\n```\n","DeepSelect 是一个针对稀疏注意力与采样场景优化的高性能 TopK 算子 CUDA 实现，专为 DeepSeek V3.2\u002FV4 系列模型中的 Sparse Attention（DSA）和 logits 采样设计。它支持 bfloat16（Lightning Indexer 场景）和 float32（Sampling 场景）输入，在小 TopK（≤4096）、大词汇表（最高 128K+）和可变批大小下，相比 PyTorch 原生 topk 实现获得 2–20 倍加速。核心特点包括内存带宽优化、可选跳过值输出、索引排序开关及严格对齐的内存布局要求。适用于大语言模型推理中的稀疏注意力计算与高效 token 采样等低延迟、高吞吐关键路径。",2,"2026-09-11 02:30:08","CREATED_QUERY"]