据llama.cpp b12207发布说明,ggml-cuda为SSM_SCAN内核新增了d_state=64快速路径,覆盖Zamba2及类似架构在循环扫描步骤中使用的64维状态。此次改动沿用已有实现模式,增加第四个模板实例,没有改变底层数学计算。
此前,GPU内核的快速路径只覆盖96、128和256三种状态大小,64维状态不在其中。发布说明称,这导致相关架构的每个扫描步骤都回退到CPU,生成速度受限于6至7 tok/s;同一硬件上的其他架构则可达到50至70 tok/s。
就本次改动的实测结果,发布说明给出的Zamba2-7B-Instruct吞吐量由11升至15.1 tok/s,提升约35%。正确性验证使用已有的CPU参考测试框架,并新增3个基于Zamba2真实超参数的测试用例,28个用例全部通过。
这次补齐快速路径后,性能仍未达到其他架构的水平。发布说明指出,另有约28%的扫描调用仍会回退到CPU,属于一个与本次改动无关、原因尚未查明的问题。新增路径解决的是64维状态的快速路径覆盖,并未消除所有CPU回退。
原始来源:b12207: ggml-cuda : add d_state=64 fast path to SSM_SCAN kernel (#31818)
