跳到主要内容
TileLang 系统教程

tilelang-tutorial.huandu.me

第 37 章 DeepSelect:算法先行的 TopK

/code/37-deepselect-topk/

03-complexity-tables/README.md1702 字节
# 示例 03:复杂度与代价模型的小表

## 目的

把报告里三个复杂度表达式
(`.../DeepSelect/docs/DeepSelect-deep-dive.zh.md:86-100`)与 `torch.topk` 的粗略代价
模型画成表,直观回答两件事:

1. 额外工作量 `E[W]`(以及它的渐近形式 `O(k log(N/k))`)**相对 `N` 有多小**,
   以及它随 `N` 的变化趋势;
2. `B2` 取多大合适——太小则 compact 频繁(`1 + k/B2` 变大),太大则缓冲区与单次
   compact 变大(`L` 变大)。

## 运行前提

* Python 3.9+,**仅标准库**。

## 运行命令

```bash
python3 complexity_tables.py
python3 complexity_tables.py --k 512 --b 1024 --b2 1024
```

## 期望输出(节选)

```
k=512  B=1024  B2=1024   L = k+B+B2 = 2560

Table 1 - the extra TopK work shrinks, relative to the input, as N grows
         N   m=N/B      H_m    E[W] bound    bound/N
     4,096       4     2.08         8,000     1.9531
    16,384      16     3.38        12,982     0.7924
    65,536      64     4.74        18,217     0.2780
   262,144     256     6.12        23,517     0.0897
 1,048,576   1,024     7.51        28,835     0.0275
16,777,216  16,384    10.28        39,480     0.0024

Table 3 - B2 trades compact frequency against compact size
N = 1,048,576, k = 512
      B2   1 + k/B2       L    E[W] bound    bound/N
     128       5.00   1,664        62,476     0.0596
     512       2.00   2,048        30,758     0.0293
   1,024       1.50   2,560        28,835     0.0275
   4,096       1.12   5,632        47,578     0.0454
```

## 怎么读

* `Table 1` 说明 `E[W]` 只随 `log N` 增长:`N` 涨 4096 倍,界只从 8,000 涨到 39,480。
* `Table 3` 里 `B2` 从 128 涨到 8192,界先降后升,**最小值落在 `B2 ≈ k` 附近**,
  与报告「对 `k=512` 取 `B=B2=1024` 比较合理」的说法一致。
* `Table 4` 是**故意粗糙**的 `torch.topk` 代价模型(读一遍 `N` + 约 `N·log2(k)` 的
  选择/排序开销),用来说明「相对于 DeepSelect 的额外工作量,`torch.topk` 多出来的
  那一截为什么随 `N` 放大」。它不是实测,也不是逐位精确的排序代价。

## 验证级别

**CPU 逻辑验证**:实际运行过。表里全是**公式算出来的数**,不是测量值;
没有任何一条性能数字来自真实硬件。

← 浏览本章全部代码