企业官网建设流程全解析

【Bug已解决】[CUDA] Attention split heuristics unsafe at small SM counts (MPS partitions): reachable div-by-zero SIGFPE in lean_attention, degenerate flash splits 解决方案

一、现象长什么样

在 macOS 上通过 Metal Performance Shaders（MPS）分区出的 CUDA 设备（或任何“SM 数量很少”的 GPU 分区）上跑 ONNX Runtime 的lean_attention内核时，注意力 split 启发式（heuristics）算出0 个 split，于是后续用 split 数做除法时触发SIGFPE（除零崩溃），或者 split 数退化成 1 但每个 split 的序列长度超过内核限制，导致输出错误。现象：

# 现象 A：直接崩溃 # Floating point exception: 8 (SIGFPE) at lean_attention.cu: split_size = S / num_splits # 因为 num_splits 被启发式算成了 0 # 现象 B：不崩溃但输出错 # num_splits 退化成 1，单个 split 的序列长度超过 kernel 最大， # attention 结果静默错误（尤其是长序列） # 现象 C：只在“小 SM 数”设备触发 # 完整 GPU（SM 多）正常；MPS 分区 / 低配卡（SM 少）才暴露 # —— 因为启发式用 num_SMs 下整除，小 SM 数时商为 0

最坑的是现象 A：进程直接 SIGFPE 崩，且只在特定设备（MPS 分区）出现，本地完整卡永远测不出，CI 也很难覆盖这种“小 SM”环境。

二、背景

lean_attention是 ONNX Runtime CUDA EP 里一种省显存、按序列维切分的注意力实现。它会根据设备 SM 数量和一个“每 SM 最优 token 数”启发式，算出把序列S切成几份（num_splits），每份在独立 stream 上算再合并。核心公式大致是：

int num_splits = (total_tokens + tokens_per_sm * num_SMs - 1) / (tokens_per_sm * num_SMs); // 向上取整

当num_SMs很小（MPS 分区可能只暴露几个 SM）且tokens_per_sm较大时，tokens_per_sm * num_SMs可能≥ total_tokens但公式里分母仍可能算出 0（如果启发式另有一处直接用num_SMs做整除、没做max(1, ...)保护），于是S / num_splits中 num_splits=0 → SIGFPE。或 num_splits 被 clamp 到 1 但内核假设 split 数 ≥ 2。

这是 CUDA 内核启发式审查里典型的坑：用设备硬件参数（SM 数）做整数除法/切分，没对“小值/零值”做保护。

三、根因

1. num_splits可能为 0：启发式用num_SMs参与整除，小 SM 数时商为 0，后续S / num_splits直接 SIGFPE（现象 A）。

2. 退化 split 超过内核上限：num_splits被 clamp 到 1，但单 split 序列长度超内核允许的max_split_tokens，输出静默错（现象 B）。

3. 缺少对小 SM 设备的测试：CI 只用完整 GPU，MPS 分区的小 SM 场景从未覆盖，除零长期存在。

本质：是注意力 split 启发式用硬件 SM 数做切分却无max(1, ...)保护与上限校验，在少 SM 设备上除零/退化，且缺测试覆盖。

四、最小可运行复现

下面用 Python 模拟“小 SM 数时 num_splits 算成 0，导致 S/num_splits 除零”：

def compute_splits_buggy(S, num_SMs, tokens_per_sm=1024): """buggy: 直接用 num_SMs 整除，小 SM 数商为 0。""" denom = tokens_per_sm * num_SMs if denom == 0: num_splits = 0 else: num_splits = (S + denom - 1) // denom return num_splits def run_attention_buggy(S, num_SMs): num_splits = compute_splits_buggy(S, num_SMs) split_size = S // num_splits # ← 若 num_splits==0 直接 ZeroDivisionError return split_size # MPS 分区只暴露 2 个 SM，tokens_per_sm 较大 try: run_attention_buggy(S=4096, num_SMs=2) except ZeroDivisionError: print("REPRO A -> SIGFPE equivalent: num_splits == 0") # 修复：保护 num_splits >= 1 def compute_splits_fixed(S, num_SMs, tokens_per_sm=1024, max_split_tokens=8192): denom = max(1, tokens_per_sm * max(1, num_SMs)) num_splits = max(1, (S + denom - 1) // denom) # 若单 split 超限，再增加 split 数 while S / num_splits > max_split_tokens: num_splits += 1 return num_splits print("fixed num_splits:", compute_splits_fixed(4096, 2)) # >=1 且合理

buggy在num_SMs=2时num_splits=0触发除零；fixed保证 ≥1 且不超限。

五、解决方案（第一层：最小直接修复）

最小修复：切分计算对num_SMs和num_splits都做max(1, ...)保护，并校验单 split 不超限：

// 修正后的 split 启发式 int num_SMs = max(1, device_props.multiProcessorCount); // 防 0 int tokens_per_sm = 1024; int denom = tokens_per_sm * num_SMs; int num_splits = max(1, (total_tokens + denom - 1) / denom); // 防单 split 超内核上限 while (static_cast<size_t>(total_tokens) / num_splits > kMaxSplitTokens) { ++num_splits; } int split_size = total_tokens / num_splits; // 现在 num_splits >= 1，安全

这一层改动最小：两处max(1, ...)+ 上限 while，除零和退化都消失。但依赖“每处启发式都加保护”，下看第二层。

六、解决方案（第二层：结构性改进）

把“注意力 split 启发式的安全约束（≥1、不超限、与 SM 数解耦）”固化成单一事实来源。下面这个 dataclass 集中管理 split 计算契约，CUDA C++ 侧和 Python 校验侧共享同一规则：

from dataclasses import dataclass, field from typing import Optional @dataclass class LeanAttnSplitPolicy: """单一事实来源：lean_attention 的 split 启发式安全约束。""" tokens_per_sm: int = 1024 max_split_tokens: int = 8192 def compute_splits(self, S: int, num_SMs: int) -> int: # 与硬件参数解耦的、永远安全的 split 计算 sms = max(1, num_SMs) denom = max(1, self.tokens_per_sm * sms) num_splits = max(1, (S + denom - 1) // denom) while S / num_splits > self.max_split_tokens: num_splits += 1 return num_splits def split_size(self, S: int, num_SMs: int) -> int: n = self.compute_splits(S, num_SMs) if n < 1: raise ValueError(f"num_splits must be >=1, got {n}") return S // n def assert_safe(self, S: int, num_SMs: int) -> None: n = self.compute_splits(S, num_SMs) if n < 1: raise AssertionError("num_splits == 0 -> SIGFPE risk") if S / n > self.max_split_tokens: raise AssertionError("single split exceeds kernel max")

这一层的关键收益：

• 永远安全：max(1, ...)保证 num_splits≥1，杜绝 SIGFPE；
• 上限校验：单 split 超内核上限时自动加 split，杜绝退化错；
• 与硬件解耦：SM 数只作启发式输入，不再能导致除零；
• 单一事实来源：所有 split 约束收口在LeanAttnSplitPolicy。

七、解决方案（第三层：断言 / CI 守护）

把第二层钉成 pytest，挂进 CI，覆盖小 SM 场景：

import pytest from your_package.lean_attn_split import LeanAttnSplitPolicy def test_no_zero_splits_at_small_sms(): # 断言 1：极小 SM 数（MPS 分区）下 num_splits 不会为 0 p = LeanAttnSplitPolicy() for sms in (0, 1, 2, 4): n = p.compute_splits(4096, sms) assert n >= 1 def test_split_size_no_div_by_zero(): # 断言 2：split_size 计算永不除零 p = LeanAttnSplitPolicy() for sms in (0, 1, 2): assert p.split_size(4096, sms) > 0 def test_single_split_within_kernel_limit(): # 断言 3：单 split 序列长度不超内核上限 p = LeanAttnSplitPolicy(max_split_tokens=8192) p.assert_safe(65536, 2) # 会自动加到足够 split 数 n = p.compute_splits(65536, 2) assert 65536 / n <= 8192 def test_full_gpu_still_works(): # 断言 4：完整 GPU（大 SM 数）行为正常 p = LeanAttnSplitPolicy() assert p.compute_splits(4096, 80) >= 1

四条断言从“小 SM 无零 split”“split_size 不除零”“单 split 不超限”“完整 GPU 正常”四面把回归钉死在 CI。

八、排查清单

CUDA attention 在 MPS 分区/小 SM 设备上 SIGFPE 或输出错时：

1. 崩在split_size = S / num_splits？num_splits必为 0，查启发式是否用num_SMs整除且无保护。
2. 不崩但输出错？num_splits退化成 1 且单 split 超内核上限，查上限校验。
3. 是否只在小 SM 设备触发？是就确认num_SMs被max(1, ...)保护。
4. 用第二层LeanAttnSplitPolicy：split 计算收口，永远 ≥1 且不超限。
5. 加第三层 pytest，覆盖num_SMs ∈ {0,1,2,4}的极端小值。
6. 所有“用硬件参数做切分”的启发式都要对小值/零值做保护。

九、小结

lean_attention的 split 启发式 bug 本质是用设备 SM 数做整数切分却无max(1, ...)保护，在小 SM 数（MPS 分区）设备上算出num_splits=0，导致S/num_splits触发 SIGFPE，或退化成 1 个超限 split 致输出静默错；且 CI 只用完整 GPU，小 SM 场景从未覆盖。修复分三层——第一层对num_SMs/num_splits加max(1, ...)并校验单 split 不超限；第二层用LeanAttnSplitPolicy这个 dataclass 把 split 安全约束收口成单一事实来源，与硬件参数解耦；第三层用四条 pytest 把“小 SM 无零 split、split_size 不除零、单 split 不超限、完整 GPU 正常”钉死在 CI。核心心法：任何用硬件参数（SM 数）做切分/整除的启发式，都必须对小值/零值做max(1,...)保护并校验上限，否则小 SM 设备必崩。