执行摘要

• 一句话：优化 LTX2 分裂 RoPE Triton 内核，合并多个 head 的 launch grid
• 推荐动作：该 PR 值得精读，尤其对需要优化 Triton kernel 以利用 GPU 的开发者。核心设计决策是使用程序块合并多个 head，这是一种常见的 GPU 优化模式（减少 program 数量，增加每个 program 的工作量以更好地隐藏延迟）。自适应 warp 数量的选择也值得参考。建议在合并到主分支前，确认其他 GPU 架构（如 A100）的基准测试结果。

功能与动机

LTX-2 模型使用分裂式旋转位置编码（RoPE），原始的 Triton 内核为每个 token-head 对启动一个独立的 program，导致 GPU 上的 launch grid 碎片化严重，降低了执行效率。特别是当 head 数量较多（如 32 head）时，program 数量激增，影响吞吐量。该 PR 旨在通过合并多个 head 的计算到单个 program 中来减少 grid 大小，提升内核性能。

实现拆解

优化通过核心里面的向量化实现，将原来每个 program 处理一个 head 改为处理一块（最多 16 个）head，从而减少 GPU program launch 数量和 grid 大小。

1. 修改 Triton 内核 _ltx2_split_rotary_kernel 的 grid 和参数：

   • 在 apply_ltx2_split_rotary_emb 函数中，将原来的二维 grid (batch * seq_len, num_heads) 改为 (batch * seq_len, triton.cdiv(num_heads, block_heads))，其中 block_heads = min(16, triton.next_power_of_2(num_heads))。这样每个 program 处理 BLOCK_HEADS 个 head。
   • 新增编译时常量 BLOCK_HEADS 和相应的 num_warps 自适应逻辑（num_warps = min(8, max(1, block_heads))），替代原来固定的 num_warps=1。

2. 修改内核内部实现以支持多 head 的向量化加载/存储：

   • 将原来标量的 head 替换为向量 heads = head_block * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)。
   • 所有地址计算和掩码都扩展为二维：mask = (heads[:, None] < num_heads) & (offsets[None, :] < half_dim)，x_base 通过广播 heads[:, None] 得到 (BLOCK_HEADS, BLOCK_HALF) 形状的地址矩阵。
   • 相应地，tl.load 和 tl.store 的指针和掩码也扩展为二维（offsets[None, :] 以保持维数一致）。

3. 保留 BF16 舍入顺序以匹配 PyTorch 参考路径：

   • 内核中的算术顺序保持不变：先计算 (x_first * cos).to(tl.bfloat16).to(tl.float32)，然后与经过 FP32 累加的 -x_second * sin 相加，确保与 PyTorch 的 BF16 操作顺序一致。

4. 验证：

   • 在 H200 上对两个代表性形状进行了微基准测试，并验证了与 PyTorch BF16 参考路径的数值一致性（max_abs_diff = 0.0）。
   • 使用 2×H200 完整运行 LTX-2 模型推理，生成视频和性能 JSON 文件，确认无功能回归。

关键文件：

• python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py（模块 JIT 内核；类别 source；类型 core-logic；符号 _ltx2_split_rotary_kernel, apply_ltx2_split_rotary_emb）: 唯一修改的文件，包含所有核心优化逻辑：grid 重划分、多 head 向量化加载/存储、自适应 num_warps。

关键符号：_ltx2_split_rotary_kernel, apply_ltx2_split_rotary_emb

评论区精华

该 PR 没有产生 review 讨论（0 条 review 评论），只有一个审核员 mickqian 批准。PR 作者 BBuf 在描述中提供了详细的基准测试数据和模型验证结果，未引起争议。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 影响：该 PR 直接影响 LTX-2 扩散模型的分裂 RoPE 计算性能。在 H200 上，短序列加速 1.5x，长序列加速 1.94x，对扩散模型的端到端生成时间有可量化的提升（总耗时约 37.9 秒，优化后预计减少 1-3 秒）。由于仅修改一个内核函数，对其他模型没有影响。团队需确保其他 GPU 架构（如 A100）上的性能不会倒退。
  • 风险标记：缺少测试覆盖

关联脉络

• 暂无明显关联 PR