执行摘要

本 PR 修复了 vllm CUDA 内核中 persistent_topk 在低共享内存设备上因回退路径内存要求不满足而导致的引擎崩溃 bug。通过将失败条件分支路由至已有的通用 top_k_per_row_decode 内核，使 GB10 等设备能够支持需要长上下文稀疏注意力的模型推理。变更局部但关键，影响范围明确。

功能与动机

在 NVIDIA GB10（SM121，约 99 KB opt-in 共享内存）等设备上运行 GLM-5.3 等 DSA 模型进行长上下文推理时，csrc/libtorch_stable/topk.cu 中的 persistent_topk 内核调度会失败。失败原因是：当协作启动所需的 CTA 数超过硬件常驻容量时，原逻辑会尝试使用 FilteredTopK 作为回退，但该回退要求至少 128 KiB 共享内存（见 Issue #21 的崩溃描述）。设备无法满足此要求，触发 STD_TORCH_CHECK 失败，引擎终止。相关背景见 Issue #45317。本 PR 的目标是为该失败路径提供一个可用的替代实现。

实现拆解

1. 变更入口与条件细化：入口是 csrc/libtorch_stable/topk.cu 中的 launch_persistent_topk 函数。在原有的 needs_cooperative && total_ctas > hw_resident_cap 失败条件内，增加了对 max_smem_per_block < 128 * 1024 的判断，精确识别出“协作启动超载”且“设备共享内存不足”的子情况。
2. 引入回退内核并推导参数：对于识别出的子情况，不再抛出错误。而是包含 #include "ops.h" 以引入 top_k_per_row_decode 函数，并调用它。调用前需要根据 lengths 张量的维度正确推导 next_n：一维张量（dim() == 2 取 size(1)）对应 MTP 布局，否则为普通解码（next_n=1）。
3. 保留原有高内存路径：对于 max_smem_per_block >= 128 * 1024 的设备，原有的 FilteredTopKRaggedTransform 调用逻辑被完整保留，确保了现有行为的兼容性。
4. 测试与验证：PR 初版在 tests/kernels/test_top_k_per_row.py 新增了 test_persistent_topk_falls_back_on_low_smem_device 参数化测试。但根据主审意见，第二个提交移除了该测试。PR 作者提供了在两块独立 GB10 设备上的详细测试结果，包括单元测试通过、编译成功和端到端集成测试（GLM-5.3 模型）成功的证据。

csrc/libtorch_stable/topk.cu

核心变更文件，修改了 launch_persistent_topk 中低共享内存设备的回退逻辑，引入新依赖并改变控制流。

// csrc/libtorch_stable/topk.cu 中变更后的关键条件分支
// 当协作启动超载且设备 opt-in 共享内存不足 128 KiB 时
if (needs_cooperative && total_ctas > hw_resident_cap) {
  // 新增条件：针对低共享内存设备，路由至通用解码内核
  if (max_smem_per_block < 128 * 1024) {
    // 从 lengths 张量维度推导 next_n，兼容一维 ( 普通解码 ) 和二维 (MTP) 形状
    const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1;
    // 调用已有的通用 top-k 解码内核作为安全回退
    top_k_per_row_decode(logits, next_n, lengths, output, num_rows,
                        logits.stride(0), logits.stride(1), TopK);
    return;
  }
  // 保留原路径：设备共享内存 >= 128 KiB 时，继续使用 FilteredTopK 回退
  cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
      logits.const_data_ptr<float>(), ...);
}

评论区精华

• 主审决策：主审 LucasWilkinson 直接批准了 PR（"Thanks for the contribution, seems reasonable enough 👍"），但对测试提出了具体修改意见："nit: please remove the overly specified test"。这体现了对代码简洁性和测试必要性的关注。
• AI 工具评论：coderabbitai[bot] 基于静态分析提出了一个潜在问题：max_smem_per_block 的值可能在 DeviceGuard 之后过时，建议在获取后确保其反映当前设备的属性。作者 lucamotz 回应此评论需要先确认问题是否实际存在，展示了对 AI 建议的审慎态度。
• CI 协作：作者在 CI 遇到与本 PR 无关的失败时，主动向主审询问处理方式（"@LucasWilkinson how do you handle this here?"），体现了良好的协作沟通。

风险与影响

1. 回归风险：核心变更仅影响“协作启动超载且共享内存<128KB”这一特定失败路径。新路径使用的 top_k_per_row_decode 是已存在的、通用的解码内核，降低了引入全新逻辑的风险。主要回归点在于该内核在此特定调用上下文下的正确性。
2. 测试覆盖：虽然作者提供了详尽的手动设备测试证据，但自动化测试中针对该路径的覆盖被移除。未来可能需要补充一个更为通用的、不依赖特定设备属性模拟的测试。
3. 性能影响：回退到 top_k_per_row_decode 可能带来与 FilteredTopK 不同的性能特征，尤其是在 MTP 场景下。PR 中未包含基准测试对比。由于这是故障恢复路径，首要目标是可用性而非最优性能，此风险可接受。

关联脉络

本 PR 是解决 NVIDIA SM12x 设备（特别是 GB10）支持 DSA 类模型长上下文推理系列工作的一部分。它紧随 Issue #45317（注意力后端选择）的解决之后，处理了后续推理流程中暴露的内核调度瓶颈。PR 描述中也明确提到了与更广泛的 SM12x 启用分支（#41834）的关系，本 PR 采用了一个更小范围、独立且可立即合并的修复方案。