执行摘要

• 一句话：用 CUDA JIT 替换 PyTorch 回退，消除 CPU-GPU 同步，加速 MoE LoRA 虚拟专家路由
• 推荐动作：值得精读。展示了如何用 CUDA JIT kernel 替换 CPU-bound 回退来消除 NCCL stall，是 LoRA/MoE 性能优化的高质量实践。设计上的双路径 fallback 和精细的测试继承结构也值得借鉴。

功能与动机

原有的 _align_block_size_torch 在专家数 >1024 时触发多个 PyTorch kernel launch 和 CPU 侧的 torch.argsort，导致 GPU 流中每层出现数百微秒的 CPU stall，进而拖慢 NCCL 集体操作（其他 rank 的 AllReduce kernel 自旋等待）。移除这些 stall 能显著提升多 TP 预填充吞吐。

实现拆解

1. 新增 _align_block_size_jit 函数（python/sglang/srt/lora/triton_ops/virtual_experts.py）：实现基于 CUDA JIT 的 align_block_size，通过单次大缓冲区分配 + 切片，将 histogram、prefix-sum、expert_ids 分配和 token 散射合并为 2-3 个 kernel launch，彻底消除 CPU-GPU 同步点。
2. 扩展 CUDA JIT kernel 支持上限（python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu）：将 EXPERTS_PER_THREAD 从 4 提升到 8，使 v2 kernel 能处理最多 8192 个专家（原为 4096），满足 LoRA 虚拟专家场景（num_moe_experts * max_loras）。
3. 测试覆盖双路径（test/registered/lora/test_virtual_experts_kernels.py）：将原有的 TestAlignBlockSizeTorchSentinelBucket 重构为基类 _AlignBlockSizeSentinelBucketBase，新增 TestAlignBlockSizeJitSentinelBucket 对 JIT 路径执行相同的 sentinel/越界验证，确保两种实现行为一致。
4. 保留 torch.compile 回退：在最后一个提交中恢复 _align_block_size_torch 作为 AMD/ROCm 平台的 fallback，并设置 _align_block_size_large 按硬件条件选择 JIT 或 torch.compile 路径。

关键文件：

• python/sglang/srt/lora/triton_ops/virtual_experts.py（模块 虚拟专家路由；类别 source；类型 core-logic；符号 _align_block_size_jit, _align_block_size_large）: 核心变更文件，添加了 _align_block_size_jit 函数并修改路由逻辑，用 CUDA JIT 替代 torch.compile fallback 消除 CPU-GPU 同步。
• python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu（模块 JIT 内核；类别 source；类型 core-logic）: CUDA JIT kernel 文件，调整了专家数上限检查和 EXPERTS_PER_THREAD 调度，以支持更大的虚拟专家数。
• test/registered/lora/test_virtual_experts_kernels.py（模块 测试；类别 test；类型 test-coverage；符号 TestAlignBlockSizeTorchSentinelBucket, _AlignBlockSizeSentinelBucketBase, _align, TestAlignBlockSizeJitSentinelBucket）: 测试文件重构，为 JIT 路径添加了相同的 sentinel/越界测试，通过基类复用确保行为一致。

关键符号：_align_block_size_jit, _align_block_size_large, launch_v2

关键源码片段

python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu

CUDA JIT kernel 文件，调整了专家数上限检查和 EXPERTS_PER_THREAD 调度，以支持更大的虚拟专家数。

// 在 v2 kernel launch 中根据专家数选择 EXPERTS_PER_THREAD
if (padded_num_experts <= 2048) {
    launch_v2(std::integral_constant<int, 2>{});
} else if (padded_num_experts <= 4096) {
    launch_v2(std::integral_constant<int, 4>{});
} else {
    // 新增分支：专家数 4096~8192 时使用 EXPERTS_PER_THREAD=8
    launch_v2(std::integral_constant<int, 8>{});
}
// 同时将上限检查从 4096 提升到 8192
RuntimeCheck(num_experts <= 8192, "moe_align_block_size: num_experts must be <= 8192");

评论区精华

copilot-pull-request-reviewer[bot] 的概览评论指出该 PR 移除了 CPU 同步和重复 LoRA 逻辑，用自定义 CUDA JIT 替换了大专家 fallback，并串联了请求级 LoRA 路由元数据。Fridge003 直接批准。作者在 PR body 中补充了详细的性能归因分析（NCCL un-stalling 效果）和内存非法访问（bs≥128）的修复说明，但没有展开讨论。

• Pull request overview by copilot-pull-request-reviewer[bot] (other): Fridge003 批准了该 PR。

风险与影响

• 风险：
  1. 硬件兼容性：CUDA JIT kernel 仅支持 NVIDIA GPU；AMD/ROCm 平台自动回退到 torch.compile，性能不降级但未获得优化收益。
  2. 专家数上限：JIT kernel 硬限制为 8191 个虚拟专家（num_moe_experts * max_loras），超出时会断言失败（_align_block_size_jit 中的 assert）。
  3. 内存越界风险：之前已出现 bs≥128 时的非法内存访问，该 PR 修复了该问题，但缓冲区对齐和 vectorized 写入的边界条件仍需关注。
     - 影响：对使用 LoRA + Expert Parallelism 的预填充场景（特别是大专家数）性能提升显著；多 TP 环境下 AllReduce 等待减少 68%。对纯 decode、小专家数或 AMD 用户影响较小。团队需确保 CI 覆盖 JIT 和 torch.compile 两条路径。
     - 风险标记：CUDA only, 专家数上限 8191, 内存越界修复

关联脉络

• PR #24246 Previous LoRA optimization PR (part of the series): 该 PR 是基于 #24246 的延续，PR body 明确说明 'based on https://github.com/sgl-project/sglang/pull/24246'，共同组成 LoRA 性能优化系列。
• PR #24248 Potential intermediate PR: 从提交历史看，多次 merge main 和分支，可能涉及其他中间 PR，但未明确提及。