Eval bug: rms_norm_f32 exceeds the CUDA gridDim.y limit (65535) at n_ctx 262144

问题场景

用户在 CUDA 后端运行 llama.cpp，加载 Qwen3.8-Flash-Next（Qwen3-Flash-Next 类超长上下文模型，相关 PR #27742）并尝试跑满原生 262,144 上下文时崩溃。原始复现用 llama-bench 单序列；后续评论显示在 llama-server 并行 slot 场景下也会出现，且报告了 rms_norm_mul_f32_cuda 的调用栈。该模型会随上下文放大 nchannels（hyper-connection 数为 4，即 nchannels = n_ctx / 4），因此是暴露出该内核限制的典型场景；任何把 nchannels 推过 65,535 的架构都可能遇到同一堵墙。

报错原文

CUDA error: invalid argument
  current device: 0, in function ggml_cuda_kernel_launch at ggml/src/ggml-cuda/common.cuh:1674
  cudaGetLastError()
ggml/src/ggml-cuda/ggml-cuda.cu:107: CUDA error

[LAUNCHFAIL] err=invalid argument grid=(1,65536,1) block=(256,1,1) shmem=128
kernel=void (*)(float const*, float*, int, long, long, long, float,
                float const*, long, long, long, uint3, uint3, uint3, uint3,
                float const*, long, long, long, uint3, uint3, uint3, uint3)

并行 slot 场景下的崩溃栈片段：

.../ggml/src/ggml-cuda/ggml-cuda.cu:107: CUDA error
#7  0x0000aef6b7c72dd8 in rms_norm_mul_f32_cuda(float const*, float const*, float const*, float*, int, int, int, int, long, long, long, long, long, long, unsigned int, unsigned int, unsigned int, unsigned int, long, long, long, unsigned int, unsigned int, unsigned int, unsigned int, float, CUstream_st*) ()

原因分析

最可能的原因：ggml/src/ggml-cuda/norm.cu 中的 rms_norm_f32（fused mul + add 变体）按 const dim3 blocks_num(nrows, nchannels, nsamples); 启动内核（出现在该文件 283、307、353、423 行附近）。当 nchannels 达到 65,536 时，grid.y 超出 CUDA 的 gridDim.y 上限 65,535，于是返回 invalid argument。同步的 binbcast.cu:317 已有 if (block_nums.z > 65535 || block_nums.y > 65535) 的退化分支，但 norm.cu 缺这个保护。

边界规律：在单序列下 n_ctx / 4 > 65535 会失败，实测 n_ctx 261,888（nchannels 65,472）通过，262,144（nchannels 65,536）失败。作者后续修正：真正的变量是每 slot 上下文 n_ctx_seq 而非 server 总 n_ctx，条件应读作 n_ctx_seq / 4 > 65535（即每 slot ≥ 262,144）。-c 262144 --parallel 4 与 -c 261888 --parallel 4（n_ctx_seq 均为 65,536）可正常加载与推理。

作者已排除的假设：2^18 的 32 位溢出、n_ctx 超过模型 262,144 上限、恰好 262,144 的 off-by-one、显存耗尽（崩溃时显存 66,186 / 96,501 MiB，约 68%）。另外 lightning-indexer.cu（QSA）与 fattn-common.cuh 的 grid y/z 不随 KV 长度扩展，不涉及。并行 slot 评论中槽位仅 121,158 token 也崩溃，不符合上述每 slot 模型，可能原因尚未定位，可能需要确认 --kv-unified 是否开启。

环境排查

• 确认 llama.cpp 版本：Issue 为 250b61446（build 10707），并称截至报告时 master 上 norm.cu 的四个无保护 launch 点未变。
• 确认 CUDA toolkit 与驱动：Issue 为 toolkit 13.1.115 / driver 595.84。
• 确认操作系统与内核：Ubuntu 26.04、kernel 7.0.0-30、glibc 2.43。
• 确认 GPU 型号与拓扑：4x RTX 3090（SM 8.6）、96,501 MiB、P2P 不可用（CNS）。
• 确认模型与量化：unsloth/Qwen3.8-Flash-Next-GGUF，UD-Q2_K_XL 与 UD-IQ4_XS 均有测试记录。
• 确认 server 参数：-c、--parallel、是否 --kv-unified，并区分总 n_ctx 与每 slot n_ctx_seq。
• 确认是否 -fa 1、-mmp 0、-ngl 999、-ot 'per_layer_token_embd=CPU' 等复现用参数。
• Issue 未提供 Python / PyTorch 依赖信息，无需据此排查。

解决步骤

1. 先确认不是显存问题：崩溃时查看显存占用。Issue 中崩溃点显存约 68%，且报错为 invalid argument 而非 OOM。
2. 开启 llama.cpp 的详细日志（如 -v），观察实际生效的 n_ctx 与启动配置，确认是否达到失败边界。
3. 计算当前 nchannels：该模型下为 n_ctx_seq / 4。在 llama-server 非 unified KV 模式下，n_ctx_seq = n_ctx / n_seq_max 并做 256 对齐（见 llama-context.cpp:290-294）。
4. 可优先尝试把每 slot 上下文降到阈值以下：单序列实测 261,888 通过、262,144 失败，可先降到约 261,888 或更低验证是否消失。
5. 若使用 llama-server 并行，注意总 -c 不直接决定是否触发，关键是每 slot n_ctx_seq；同时留意长上下文下 compute buffer 增长远快于 KV（4K→262K：KV +8.1 GiB，compute +29.6 GiB），可能先撞显存。
6. 若要根修，方向是给 norm.cu 的四个 launch 点补上类似 binbcast.cu:317 的 block_nums.y > 65535 退化分支，或把 nchannels 拆分到 grid.y / grid.z。Issue 作者表示若 maintainer 确认首选形态愿提 PR，但当时未有确认，故不要视为已可用修复。

验证方法

用相同模型与参数复跑失败配置（例如单序列 -c 262144），确认不再出现 rms_norm_f32 / rms_norm_mul_f32_cuda 的 invalid argument；同时对照 -c 261888 等已知通过配置，确认降级后能正常加载并推理。若采用 patch，需验证四个 launch 点在 nchannels 超过 65,535 时走退化分支且结果数值正确。Issue 未提供官方修复版本，验证只能针对本地改动或参数调整。

参考来源

ggml-org/llama.cpp #27901

AI 工具推荐

想把多个 AI 模型放在一个入口？

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型，适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接：通过此链接购买，我可能获得佣金，不影响你的价格。