执行摘要

• 一句话：修复 MoE LoRA 对齐内核 grid 越界导致 CUDA 非法访问
• 推荐动作：建议阅读。该 PR 展示了如何诊断 CUDA kernel 中因 grid 大小不足导致的 off-by-one 错误，并采用防御性 guard 增强健壮性。测试设计中使用 sentinel 值检测未初始化输出的思路值得借鉴。对于维护 MoE LoRA 相关代码的工程师，此修复直接解决了常见的 illegal address 崩溃。

功能与动机

当 max_loras 等于实际活跃 LoRA 数量且 batch 包含 base 模型 token 时，active_lora_ids 数组长度为 max_loras + 1（包含 -1），但 grid 只启动 max_loras 个 block，导致最后一个 LoRA slot 的 token 被跳过，输出缓冲区保持未初始化，触发 cudaErrorIllegalAddress。此问题与 #32235 描述的 Triton fused_moe_lora_kernel 类似，但发生在 C++ align kernel 上。PR body 明确说明了触发条件和修复思路。

实现拆解

1. 在 csrc/moe/moe_align_sum_kernels.cu 的 moe_lora_align_block_size 函数中，将启动三个对齐内核的 grid 大小从 max_loras 改为 max_loras + 1，使得额外增加的 base-model slot（lora_id = -1）也能被处理。
2. 在每个内核的开头添加 lora_id >= max_loras 检查，与已有的 lora_id == -1 和 adapter_enabled[lora_id] == 0 组合为统一的跳过条件，防止伪造或意外的 lora_id 导致越界读写。
3. 为 lora_count_and_sort_expert_tokens_kernel 新增 adapter_enabled 参数传递和检查，使得禁用 adapter 的 slot 也被跳过，避免读取未初始化的 token_mask 造成 sorted_token_ids 污染。
4. 在 tests/lora/test_moe_lora_align_sum.py 中新增辅助函数 _build_and_run_align 和三个回归测试，分别覆盖混合 base+LoRA 批次（验证最后一个 LoRA slot 被正确填充）、禁用 adapter slot（验证输出缓冲区保持 sentinel），以及 lora_id 越界（验证 guard 生效）。测试使用 sentinel 值预填充输出缓冲区，断言内核写入预期范围。

关键文件：

• csrc/moe/moe_align_sum_kernels.cu（模块 MoE内核；类别 source；类型 core-logic；符号 moe_lora_align_block_size, moe_lora_align_block_size_kernel, lora_count_and_sort_expert_tokens_kernel, moe_lora_align_block_size_small_batch_expert_kernel）: 修复 MoE LoRA 对齐内核的 grid 大小和边界检查，是 bug 根因所在
• tests/lora/test_moe_lora_align_sum.py（模块 LoRA测试；类别 test；类型 test-coverage；符号 _build_and_run_align, test_moe_lora_align_block_size_mixed_base_and_lora, test_moe_lora_align_block_size_disabled_adapter_untouched, test_moe_lora_align_block_size_lora_id_oob_guard）: 添加回归测试，覆盖混合 base+LoRA、禁用 adapter、越界 lora_id 等场景，确保内核修复的有效性

关键符号：moe_lora_align_block_size, moe_lora_align_block_size_kernel, lora_count_and_sort_expert_tokens_kernel, moe_lora_align_block_size_small_batch_expert_kernel, _build_and_run_align, test_moe_lora_align_block_size_mixed_base_and_lora, test_moe_lora_align_block_size_disabled_adapter_untouched, test_moe_lora_align_block_size_lora_id_oob_guard

关键源码片段

csrc/moe/moe_align_sum_kernels.cu

修复 MoE LoRA 对齐内核的 grid 大小和边界检查，是 bug 根因所在

// 文件 : csrc/moe/moe_align_sum_kernels.cu
​
// 将 grid 从 max_loras 改为 max_loras + 1，确保 base-model slot 也被覆盖
// active_lora_ids 长度为 max_loras + 1，其中索引 0 可能为 -1（base 模型）
kernel<<<max_loras + 1, blockDim, shared_mem, stream>>>( 
    // ... 参数 ...
);
​
// 每个对齐内核中添加防御性 guard，以下以 lora_count_and_sort_expert_tokens_kernel 为例
// 该内核在 review 中被指出缺少 adapter_enabled 检查，现已修复
__global__ void lora_count_and_sort_expert_tokens_kernel(
    // ... 原有参数 ...,
    int32_t max_loras,
    int32_t* lora_ids,
    int32_t* adapter_enabled, // 新增参数，用于检查 adapter 是否启用
    bool has_expert_map) {
  int lora_idx = blockIdx.x;
  int lora_id = lora_ids[lora_idx];
  // 跳过 base 模型（-1）、越界 lora_id、禁用 adapter
  if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) {
    return; // 不执行后续排序逻辑，避免读写未初始化数据
  }
  // ... 原有排序业务逻辑 ...
}

评论区精华

• gemini-code-assist[bot] 指出 lora_count_and_sort_expert_tokens_kernel 未检查 adapter_enabled，可能读取未初始化的 token_mask 并导致 sorted_token_ids 污染。作者根据建议添加了 adapter_enabled 参数传递和检查，并补充了对应测试。

• jeejeelee 询问 lora_id >= max_loras 检查是否必要，因为正常路径中 lora_id 应该始终小于 max_loras。作者解释这是防御性编程，防止伪造的 lora_ids 导致越界，并引用现有测试中使用 lora_ids = torch.arange(max_loras + 2) 的场景（该测试在 grid 扩大后会启动额外 block，需要 guard 保护）。

• sort kernel 缺少 adapter_enabled 检查 (correctness): 作者根据建议添加了 adapter_enabled 参数传递和检查，并补充了对应的单元测试来覆盖禁用 adapter 场景。

• 防御性 guard lora_id >= max_loras 的必要性 (design): 作者解释这是防御性编程，防止伪造或异常的 lora_ids 导致越界，并引用现有测试中使用 lora_ids = torch.arange(max_loras+2) 时 grid 扩大后会启动额外 block，需要该 guard 保护。

风险与影响

• 风险：主要风险点：
  • grid 增大一个 block 对性能影响可忽略，且由于 guard 会立即返回。
  • 防御性 guard 在正常路径下不触发，不会改变行为。
  • 若 adapter_enabled 数组长度不足 max_loras + 1（生产环境中分配长度为 max_loras + 1），或 lora_ids 未按预期排序，但现有测试已覆盖多种边界。
  • 测试覆盖了混合批次、禁用 adapter 和越界场景，降低了回归风险。
  • 影响：直接影响使用 MoE LoRA 特性的用户，修复了在 mixed batch（base 模型 token 与 LoRA token 混用）且 max_loras 等于活跃 LoRA 数时出现的 CUDA 非法内存访问错误，显著提升稳定性。对不使用 MoE 或 LoRA 的用户无影响。改动范围仅两个文件，兼容性保持不变。测试套件的增强也提高了后续修改的安全性。
  • 风险标记：混合 batch 场景, 网格大小调整, 防御性 guard

关联脉络

• PR #32277 [BugFix] fix fused_moe_lora launch grid bugs: 本PR修复的C++内核grid问题与#32277针对Triton内核的修复是同一类off-by-one错误，本PR借鉴了其grid增大思路并扩展到了C++ align内核。