【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 ShadersMPS分区出的 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 数”设备触发 # 完整 GPUSM 多正常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_splits0 → SIGFPE。或 num_splits 被 clamp 到 1 但内核假设 split 数 ≥ 2。这是 CUDA 内核启发式审查里典型的坑用设备硬件参数SM 数做整数除法/切分没对“小值/零值”做保护。三、根因num_splits可能为 0启发式用num_SMs参与整除小 SM 数时商为 0后续S / num_splits直接 SIGFPE现象 A。退化 split 超过内核上限num_splits被 clamp 到 1但单 split 序列长度超内核允许的max_split_tokens输出静默错现象 B。缺少对小 SM 设备的测试CI 只用完整 GPUMPS 分区的小 SM 场景从未覆盖除零长期存在。本质是注意力 split 启发式用硬件 SM 数做切分却无max(1, ...)保护与上限校验在少 SM 设备上除零/退化且缺测试覆盖。四、最小可运行复现下面用 Python 模拟“小 SM 数时 num_splits 算成 0导致 S/num_splits 除零”def compute_splits_buggy(S, num_SMs, tokens_per_sm1024): 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_splits0 直接 ZeroDivisionError return split_size # MPS 分区只暴露 2 个 SMtokens_per_sm 较大 try: run_attention_buggy(S4096, num_SMs2) except ZeroDivisionError: print(REPRO A - SIGFPE equivalent: num_splits 0) # 修复保护 num_splits 1 def compute_splits_fixed(S, num_SMs, tokens_per_sm1024, max_split_tokens8192): 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_SMs2时num_splits0触发除零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_castsize_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(fnum_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(): # 断言 2split_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_tokens8192) 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 或输出错时崩在split_size S / num_splitsnum_splits必为 0查启发式是否用num_SMs整除且无保护。不崩但输出错num_splits退化成 1 且单 split 超内核上限查上限校验。是否只在小 SM 设备触发是就确认num_SMs被max(1, ...)保护。用第二层LeanAttnSplitPolicysplit 计算收口永远 ≥1 且不超限。加第三层 pytest覆盖num_SMs ∈ {0,1,2,4}的极端小值。所有“用硬件参数做切分”的启发式都要对小值/零值做保护。九、小结lean_attention的 split 启发式 bug 本质是用设备 SM 数做整数切分却无max(1, ...)保护在小 SM 数MPS 分区设备上算出num_splits0导致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 设备必崩。