llama.cpp b11405:CUDA lightning indexer 分块优化
原标题:b11405
AI 摘要
llama.cpp 发布 b11405 版本,主要针对 CUDA 后端优化 lightning indexer 内核。改动包括按 keys 和 tokens 对 4 个 head 进行分块、在 float 精度下计算乘积以避免 half2 溢出、每个 key 只做一次精度扩展供所有 head 使用,并修复 ROCm 上的寄存器溢出问题。优化后该内核在 R9700 上提速 36 倍,CUDA 上略有提升,同时提供覆盖 macOS、Linux、Windows、Android 的多平台预编译包。
正文节选
<details open> cuda: tile the lightning indexer over keys and tokens for 4 heads (#29901) * cuda: tile the lightning indexer over keys and tokens for 4 heads With too few heads for a wmma tile, a block scores 64 keys against 8 tokens: the keys are staged once in half precision, the queries one head at a time, and each thread owns one key for two tokens, so no dot product needs a cross thread reduction. Batches smaller than a token tile keep the vector kernel. test-backend-ops measures 4 heads