Jane Street 用 NVIDIA cuTile 实现 Mamba-2 高性能前向核
blelbach · x · 2026-10-23
Jane Street 与 NVIDIA 联合发布技术博客,介绍用 cuTile(基于 Python 的 tile 级 DSL)实现 Mamba-2 状态空间模型的前向传播核函数。要点:
- 实现并优化了三个自定义 kernel:chunkcumsum(log 空间累积衰减项)、fwdpassstate(用 ct.mma 走 Tensor Core 推进隐藏状态)、computeout(把状态项、衰减掩码交互矩阵和 skip connection 融合进单个 kernel)。
- 优化手段包括 DRAM 延迟提示、操作数交换路由到 Tensor Memory、适配 TMA 的转置内存布局,以及按输入形状自动搜索块大小的 autotuner。
- 展示了 cuTile 用高层 tile 抽象自动对接 Tensor Core 等硬件特性的能力,工程师无需手写线程/warp 级调度。
「Infra」频道最新
- M3 Max 36GB 五引擎推理实测:Kapitsa 密集 41 tok/s 居首 — Hyungsun · 2026-10-24
- Ramjet 替换 nginx 使 domyn swarm 多 agent 会话提速最高 61% — DoggoProfessor959 · 2026-10-24
- 美光 G11 NAND:同面积容量翻倍,2027 下半年量产 — rohanpaul_ai · 2026-10-24
- PSA:从 SSD 流式读取 n-gram 表会实打实损耗 SSD 寿命 — Chromix_ · 2026-10-24
- ISTA 发布 DeepSeek-V4.1-Flash 极限量化版:2 比特性能不掉,单卡 12GB 可跑 763B 模型 — Loginhe · 2026-10-24
- Nebius Token Factory 上线缓存 Token:命中前缀自动按缓存价计费 — HowDevelop · 2026-10-24