feat: execute FlashKDA on RTX 5090
This commit is contained in:
wuyang
@@ -22,7 +22,7 @@
|
||||
- 没有在单张 RTX 5090 上加载 1.56 TB checkpoint；
|
||||
- 没有获得真实 token hidden states、线上 expert load 或生产 cache trace；
|
||||
- 没有把随机向量上的 router 行为写成真实数据分布；
|
||||
- 没有把未成功执行的 FlashKDA kernel 写成“本机 benchmark”；
|
||||
- 没有把合法 synthetic shape 上的 FlashKDA 实测写成“K3 checkpoint 已跑通”；
|
||||
- 没有替官方解释下面发现的 `A_log` 形状不一致。
|
||||
|
||||
## 2. 一手工件与校验
|
||||
@@ -268,7 +268,7 @@ RMS=1 的各向同性随机 hidden vectors，再比较 top-16：
|
||||
|
||||
### 7.1 官方仓库事实
|
||||
|
||||
FlashKDA `1ce47ea3`：
|
||||
[FlashKDA 官方仓库](https://github.com/MoonshotAI/FlashKDA) `1ce47ea3`：
|
||||
|
||||
- CUTLASS kernels；
|
||||
- 支持 `90a / 100a / 103a / 120a`；
|
||||
@@ -288,51 +288,127 @@ FlashKDA `1ce47ea3`：
|
||||
|
||||
这些是作者仓库 benchmark，不是本站复跑值。
|
||||
|
||||
### 7.2 本机真实构建边界
|
||||
### 7.2 从主机构建失败到隔离 wheel
|
||||
|
||||
本机：
|
||||
GPU 主机：
|
||||
|
||||
- RTX 5090，compute capability `12.0`；
|
||||
- PyTorch `2.11.0+cu128`；
|
||||
- PyTorch CUDA `12.8`；
|
||||
- Ubuntu 26.04，glibc `2.43`；
|
||||
- 常用环境为 PyTorch `2.11.0+cu128` / CUDA `12.8`；
|
||||
- 官方源码明确包含 `sm_120a`，所以不是 GPU architecture 缺失。
|
||||
|
||||
两次可复现构建：
|
||||
主机路径依次暴露了两个不同问题：
|
||||
|
||||
1. 系统 `g++ 15.2`：PyTorch extension 在编译前拒绝，CUDA 12.8 要求 host compiler `<14`；
|
||||
2. 临时解包 `g++ 13.4`：成功进入 nvcc，但 CUDA 12.8 headers 与当前 glibc math declarations
|
||||
在 `cospi / sinpi / rsqrt` exception specification 处冲突。
|
||||
冲突；
|
||||
3. 新建 PyTorch `2.11.0+cu130` 环境并改用 CUDA 13.1 后，版本主线已经匹配，但仍在
|
||||
`rsqrt / rsqrtf` exception specification 处失败。
|
||||
|
||||
结论：
|
||||
第三步说明“升级到 CUDA 13”还不够。[CUDA 13.1 官方 Linux support matrix](https://docs.nvidia.com/cuda/archive/13.1.0/cuda-installation-guide-linux/index.html) 当前列出的
|
||||
glibc 上限是 2.41，而主机是 2.43；NVIDIA 官方论坛也记录了同类 CUDA 13.1 / 新 glibc
|
||||
数学声明[冲突案例](https://forums.developer.nvidia.com/t/fedora-43-and-nvcc-cuda13-1-error-exception-specification-is-incompatible-rsqrt-rsqrtf/354510)。这里能写的是“构建失败与超出验证矩阵一致”，不能写成 CUDA 对未来 glibc
|
||||
永久不兼容。
|
||||
|
||||
- kernel 尚未在本站机器执行；
|
||||
- 失败与 README 的 CUDA 12.9+ 要求一致；
|
||||
- 不能把 `sm_120a` 支持写成本机已经跑通；
|
||||
- 下一次应使用匹配 PyTorch 的 CUDA 12.9+ toolchain 或官方容器后再复跑 correctness + benchmark。
|
||||
最终构建固定为：
|
||||
|
||||
| layer | pinned value |
|
||||
|---|---|
|
||||
| base image | `nvidia/cuda:13.0.2-devel-ubuntu24.04` |
|
||||
| container libc | glibc `2.39` |
|
||||
| Python | `3.12` |
|
||||
| PyTorch | `2.11.0+cu130` |
|
||||
| target | `sm_120a` |
|
||||
| FlashKDA | `1ce47ea3bb22c84eb9cc665028399cf35e8ffb0b` |
|
||||
| wheel | `flash_kda-0.0.1+1ce47ea-cp312-cp312-linux_x86_64.whl` |
|
||||
| wheel SHA-256 | `14687b6d84a256d4552f0c73ccf93a601be582aeabcdf49ae3a409266872158d` |
|
||||
|
||||
容器没有 NVIDIA runtime，也不需要 GPU 才能为固定 architecture 编译。CUDA 编译器实际产出
|
||||
`sm_120a` 对象，最终 wheel 为 3,806,174 bytes。wheel 随后装入主机独立 CPython 3.12.11 /
|
||||
PyTorch 2.11.0+cu130 环境；该运行环境能看到 RTX 5090、capability 12.0，并成功导入
|
||||
`flash_kda` 与 `flash_kda_C`。
|
||||
|
||||
### 7.3 exact correctness：输出与 recurrent state 都必须相等
|
||||
|
||||
正确性对照直接调用官方仓库的 `tests/torch_ref.py`，而不是本站重写一份近似参考。
|
||||
测试输入满足公开 API 的 `A_log[H]`；它们故意不使用 checkpoint 的 `A_log[128]`，因此不会
|
||||
偷偷替未决形状冲突下结论。
|
||||
|
||||
| case | sequence lengths | H | output exact | final state exact | max abs error |
|
||||
|---|---|---:|---:|---:|---:|
|
||||
| one chunk | `[16]` | 1 | yes | yes | 0 |
|
||||
| partial tail | `[17]` | 1 | yes | yes | 0 |
|
||||
| two chunks | `[32]` | 1 | yes | yes | 0 |
|
||||
| multi-chunk / multi-head | `[65]` | 2 | yes | yes | 0 |
|
||||
| K3 head count | `[17]` | 96 | yes | yes | 0 |
|
||||
| varlen partial chunks | `[17,31]` | 2 | yes | yes | 0 |
|
||||
|
||||
六组用 BF16 q/k/v/g、BF16 beta logits、F32 `A_log/dt_bias` 和 BF16 recurrent state；
|
||||
output 与 final state 全部逐元素相等，所有 max / mean absolute difference 都是 0。
|
||||
这把证据从“源码声称支持 120a”推进成了“本站 RTX 5090 实际执行并通过官方参考”。
|
||||
|
||||
### 7.4 本机受控计时：只报本站 FlashKDA，不伪造本机 FLA 对照
|
||||
|
||||
协议：
|
||||
|
||||
- deterministic synthetic tensors；
|
||||
- `D=128`，K3 shape 使用 `H=96, T_total=8192`；
|
||||
- warmup 20；
|
||||
- 每轮 100 次、3 repeats，即每个 state mode 300 个 CUDA Event samples；
|
||||
- output buffer 重用；
|
||||
- 分开测 BF16 state、无 state I/O、FP32 state；
|
||||
- throughput 是 sequence tokens/s，不是 head-tokens/s；
|
||||
- 没有安装并复跑 FLA，所以不计算本站 speedup。
|
||||
|
||||
| RTX 5090 case | state mode | mean | P95 | sequence tokens/s |
|
||||
|---|---|---:|---:|---:|
|
||||
| fixed `[8192]` | BF16 | 2.6210 ms | 2.6437 ms | 3.126 M/s |
|
||||
| fixed `[8192]` | none | 2.6154 ms | 2.6362 ms | 3.132 M/s |
|
||||
| fixed `[8192]` | FP32 | 2.6267 ms | 2.6500 ms | 3.119 M/s |
|
||||
| varlen `[1300,547,2048,963,271,3063]` | BF16 | 2.3335 ms | 2.3574 ms | 3.511 M/s |
|
||||
| same varlen | none | 2.3104 ms | 2.3254 ms | 3.546 M/s |
|
||||
| same varlen | FP32 | 2.3596 ms | 2.3845 ms | 3.472 M/s |
|
||||
|
||||
fixed case 三种 mode 合并的 peak allocated 为 1,755.7 MiB；varlen 为 1,845.7 MiB。
|
||||
作者 H20 fixed 的 `2.6220 ms` 与本站 BF16 fixed 的 `2.6210 ms` 数值接近只是一次观测，
|
||||
不能据此宣布 RTX 5090 与 H20 等价：软件栈、频率、功耗、输入、测量时段和对照 kernel
|
||||
都不足以支持这种外推。
|
||||
|
||||
## 8. 可复现实验入口
|
||||
|
||||
脚本：
|
||||
checkpoint 工件脚本：
|
||||
|
||||
```text
|
||||
experiments/k3/checkpoint_probe.py
|
||||
```
|
||||
|
||||
FlashKDA 构建与执行：
|
||||
|
||||
```text
|
||||
experiments/k3/flashkda/Dockerfile
|
||||
experiments/k3/flashkda/build_reference_helper.py
|
||||
experiments/k3/flashkda/run_probe.py
|
||||
experiments/k3/flashkda/README.md
|
||||
```
|
||||
|
||||
提交的数据快照：
|
||||
|
||||
```text
|
||||
src/data/k3-artifact-snapshot.json
|
||||
src/data/k3-flashkda-runtime.json
|
||||
```
|
||||
|
||||
脚本会：
|
||||
两条实验路径共同完成：
|
||||
|
||||
1. 解析 config、index 与 selected headers；
|
||||
2. 校验小范围字节长度和 SHA-256；
|
||||
3. 统计真实 KDA / router 参数；
|
||||
4. 运行明确标注的 synthetic router counterexample；
|
||||
5. 解析 FlashKDA 官方 H20 / GB200 benchmark；
|
||||
6. 输出本机环境与构建边界；
|
||||
7. 不提交原始权重。
|
||||
6. 在隔离 CUDA 13.0 / glibc 2.39 环境编译 `sm_120a` wheel；
|
||||
7. 用官方 torch reference 跑六组 exact-match；
|
||||
8. 在 RTX 5090 跑 fixed / varlen、三种 state mode 的受控计时；
|
||||
9. 不提交原始权重或二进制 wheel，只提交 checksum、环境、代码和 JSON 结果。
|
||||
|
||||
## 9. 网站实现合同
|
||||
|
||||
@@ -341,7 +417,7 @@ src/data/k3-artifact-snapshot.json
|
||||
1. **Layer map**：93 层真实 config 条带；显示 KDA/MLA、dense/MoE、AttnRes block。
|
||||
2. **Tensor anatomy**：checkpoint / expert / vision / MLA tensor shape 与数量。
|
||||
3. **Parameter audit**：真实 Range statistics，并把 `A_log` mismatch 放在主视区。
|
||||
4. **Reproduction boundary**：官方 benchmark、本站构建失败点、synthetic router counterexample。
|
||||
4. **Reproduction boundary**：作者 benchmark、本站 exact / latency、构建链、未决形状冲突与 synthetic router counterexample。
|
||||
|
||||
每个视图必须显示证据类型：
|
||||
|
||||
@@ -359,7 +435,9 @@ src/data/k3-artifact-snapshot.json
|
||||
- [x] selected open-weight ranges 做真实参数统计；
|
||||
- [x] 发现并限定 `A_log` shape inconsistency；
|
||||
- [x] FlashKDA RTX 5090 构建尝试留下可复现边界；
|
||||
- [ ] 使用 CUDA 12.9+ 匹配环境跑 FlashKDA exact correctness；
|
||||
- [x] 使用 CUDA 13.0 / glibc 2.39 隔离环境产出可校验 `sm_120a` wheel；
|
||||
- [x] RTX 5090 上六组 FlashKDA / official torch reference exact correctness；
|
||||
- [x] K3 fixed / varlen shape、三种 state mode 的 1,800 个 CUDA Event samples；
|
||||
- [ ] 取得真实 hidden-state / router load trace；
|
||||
- [ ] 取得可加载的 reduced checkpoint、官方 trace 或多机资源；
|
||||
- [ ] 对 Figure 3 / 4 / 5 做真实数值重绘；
|
||||
|
||||
Reference in New Issue
Block a user