执行摘要

• 一句话：避免FLA L2-norm按token数重编译
• 推荐动作：该 PR 值得精读，展示了 Triton 内核性能优化中避免过度特化的经典技巧（do_not_specialize）。对于维护 Triton 内核的团队，应评估是否有其他编译时常量（如序列长度、批次大小）可转为动态参数以降低编译开销。

功能与动机

VLM 服务中，图像分辨率变化会导致视觉/文本 token 长度频繁变动，原代码将 token 总数 T 作为编译时常量，为每个独特长度编译新的 Triton cubin，造成显著的冷启动延迟与 GPU 内存浪费。PR body 明确指出 'This is especially visible for VLM serving because image resolution changes the flattened vision/token length from request to request.'，并在微基准测试中验证：6 种 token 长度下，cubin 数量从 6 降为 1。

实现拆解

1. 移除已废弃的 NB 参数：

   • 在 l2norm_fwd_kernel 函数签名中删除 NB: tl.constexpr 参数及调用处的 NB=NB 传参（l2norm_fwd 函数中）。
   • NB 原用于 Triton 编译时控制块数量，但已不被使用。

2. 将 T 从编译时常量改为运行时参数：

   • 将 T: tl.constexpr 改为 T（无 tl.constexpr）。
   • 在 @triton.jit 装饰器中添加 do_not_specialize=["T"]，指示 Triton 编译器不因 T 的不同值而生成新内核。
   • 此举确保所有 token 长度共享同一个编译后的 cubin，同时保持 D, BT, BD 等与内存布局强相关的参数仍在编译期固定。

3. 影响范围：

   • 仅修改 python/sglang/kernels/ops/attention/fla/l2norm.py 一个文件，+2/-5 行。
   • 该 l2norm_fwd 函数被 GDN 和 KDA 注意力路径共用，因此优化对两者均有效。
   • 无测试、配置或部署配套改动。

关键文件：

• python/sglang/kernels/ops/attention/fla/l2norm.py（模块 Kernel；类别 source；类型 core-logic；符号 l2norm_fwd_kernel, l2norm_fwd）: 核心变更文件，包含 L2-norm Triton 内核的优化：移除 NB 参数、将 T 从编译时常量改为运行时参数、添加 do_not_specialize。

关键符号：l2norm_fwd_kernel, l2norm_fwd

关键源码片段

python/sglang/kernels/ops/attention/fla/l2norm.py

核心变更文件，包含 L2-norm Triton 内核的优化：移除 NB 参数、将 T 从编译时常量改为运行时参数、添加 do_not_specialize。

@triton.jit(do_not_specialize=["T"]) # T 是运行时变量，不同 token 长度复用同一内核
def l2norm_fwd_kernel(
    x, y, eps,
    T, # 之前是 tl.constexpr，现在改为运行时参数
    D: tl.constexpr,
    BT: tl.constexpr,
    BD: tl.constexpr,
    # ... 内核实现保持不变
):
    pass
​
def l2norm_fwd(x, eps=1e-6):
    D = x.shape[-1]
    T = x.numel() // D
    if D <= 512:
        # NB = triton.cdiv(T, 2048) # 已废弃，删除
        def grid(meta):
            return (triton.cdiv(T, meta["BT"]),)
        l2norm_fwd_kernel[grid](
            x, y, eps,
            T=T, # 不再传递 NB
            D=D, BD=BD,
        )

评论区精华

无人工 review 评论，仅有 gemini-code-assist[bot] 的自动代码审查，确认了变更内容且未提出异议。该机器人还附带 Gemini Code Assist 服务即将下线的通知，与技术内容无关。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 数值正确性风险（低）：将 T 由编译时常量改为运行时参数，可能影响 Triton 编译器对循环展开或寄存器分配的优化。但 do_not_specialize 是 Triton 官方支持的机制，且测试已在 H200 上覆盖 9 个 token 长度（1,7,127,511,1025,2049,4097）配合 4 种特征维度（64,128,256,512），所有结果与 PyTorch FP32 参考比对通过。
  2. 回归风险（低）：改动极小（仅删去 5 行、修改装饰器参数），不涉及逻辑变更，且移除了未使用的 NB 参数，不会引入新 bug。
  3. 兼容性风险（低）：函数签名变化（删除 NB 参数、T 不再为 constexpr），任何外部直接调用 l2norm_fwd 的代码需要适配，但仓库内无其他调用点。
• 影响：
  • 性能影响（高）：实测对 Qwen3.6 多模态预填充场景 TTFT 降低 26%-40%，吞吐量提升 35%（满载时）。冷 Triton 缓存下 cubin 从 6 个减少为 1 个，首次调用延迟显著降低。
  • 用户影响（中）：所有使用 FLA L2-norm 的模型（如 GDN、KDA 注意力路径）受益，尤其 VLM 多模态服务用户。无需配置或 API 变更。
  • 系统影响（低）：无部署或配置变更。
  • 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #31109 Remove QServe and FBGEMM FP8 quantization: 同属量化路径清理与性能优化，涉及 Triton kernel 改动。