fix: prevent first-step recompilation in MiniMax-H3 attention

.claude/skills/support_model_compile/SKILL.md

@@ -55,7 +55,7 @@ block/phase 族及各自签名：
 block/phase cache key 与预期 cache 数量：
 运行时分支和 request-specific 状态：
 可能分流的第三方算子：
-正式 graph signature（shape/token/frame/dtype/layout/branch/stage）：
+正式 graph signature（原始 shape/token/frame、packing/padding/SP 后各 rank 的本地 tensor shape、dtype/layout/branch/stage）：
 warmup 是否覆盖相同 signature 和真实 leaf kernel：
 ```
 
@@ -130,7 +130,7 @@ compile 不会优化图外的 CPU→GPU 拷贝、buffer 交换或 stream synchro
 - eager 与 compile 是否数值一致；
 - 性能对照是否真的使用相同 kernel。
 
-需要隔离第三方实现时，在模块作用域定义 `torch.library.custom_op` 并注册 fake。custom op 只建立图边界，不会消除其 kernel launch，也不会跨边界融合。
+需要隔离第三方实现时，在模块作用域定义 `torch.library.custom_op` 并注册 fake。custom op 只建立 Dynamo 图边界，不会消除其 kernel launch，也不会跨边界融合。算子内部的 Triton、CUDA extension 或第三方 JIT 仍可能按运行时签名编译和特化；按 [自定义算子内层编译与 Triton 动态标量](references/implementation-patterns.md#自定义算子内层编译与-triton-动态标量) 继续检查。
 
 已确认不兼容且没有可靠 fallback 的组合应在初始化阶段明确报错，不要静默换 kernel。只有精确定位到公共算子并验证所有受影响路径后，才修改 `lightx2v/common/ops`。
 
@@ -179,7 +179,7 @@ git diff --check
 
 ## 7. 评测稳态收益
 
-两边都启用 warmup，并通过正式调用链覆盖相同的 shape、token/frame、dtype、layout、branch、stage、offload 对象生命周期和真实 leaf-op dispatch；分辨率相同不等于 graph 相同。目标模型 warmup 不可靠时，先使用 `../support_model_warmup/SKILL.md` 修复，再评测 compile。
+两边都启用 warmup，并通过正式调用链覆盖相同的 graph signature、offload 对象生命周期和真实 leaf-op dispatch。graph signature 不只包含原始 shape/token/frame、dtype、layout、branch 和 stage；多模态或序列并行路径还应比较 packing/padding/SP 后各 rank 的本地 tensor shape，尤其是实际 q/k 长度。分辨率相同不等于 graph 相同。目标模型 warmup 不可靠时，先使用 `../support_model_warmup/SKILL.md` 修复，再评测 compile。
 
 - 每个模式先完成一轮 warmup + 正式请求的功能验收；
 - 每组至少三轮；

.claude/skills/support_model_compile/references/casebook.md

@@ -20,6 +20,7 @@
 - [14. 最终输出相近但数值不一致](#14-最终输出相近但数值不一致)
 - [15. warmup 后首请求因共享 config 扩充而重编译](#15-warmup-后首请求因共享-config-扩充而重编译)
 - [16. 多模态 timestep layout 在后续 step 重编译](#16-多模态-timestep-layout-在后续-step-重编译)
+- [17. custom op 外层稳定但内层 Triton 仍冷编译](#17-custom-op-外层稳定但内层-triton-仍冷编译)
 
 ## 1. staging block 重复编译
 
@@ -145,3 +146,11 @@
 - **根因**：多模态 scheduler 的首步可能让各模态共享同一 timestep，后续 step 则产生多个 unique timesteps。以 MiniMax-H3 为例，Step 0 的 `temb.shape[0]` 为 1，后续 step 为 2；只 warmup Step 0 无法覆盖后一种图。
 - **处理**：在同一次正式 scheduler 状态上执行能覆盖每种稳定 timestep layout 的最少代表 step，并保持 `step_pre → infer → step_post`。不要默认跑完整 denoise loop，也不要仅因 shape 写成 `(H,W,T)` 就认为 graph signature 已覆盖。
 - **验收**：重编译应全部发生在服务 ready 前；再走真实 HTTP 请求，确认每个正式 step 均无新 recompile，并分别报告新增 warmup 时间和正式 E2E。
+
+## 17. custom op 外层稳定但内层 Triton 仍冷编译
+
+- **现象**：warmup 完成后，正式 Step 1 仍显著慢于稳态；custom op 已消除目标 Dynamo graph break 或 recompile，却只消除了一部分首轮耗时。
+- **确认**：比较 warmup 与正式请求在 packing、padding 和 SP 后各 rank 的实际 q/k 长度及 sparse block 数；在服务 ready 前后记录 Triton/第三方 JIT 日志或缓存产物，并用 dense/backend control 排除 runner 整体 warmup 失效。
+- **根因**：custom op 只隔离外层 Dynamo。内部 leaf kernel 仍可能按序列长度的数值或对齐关系特化；即使视频空间和时间尺寸相同，不同 prompt 经 packing 和 SP 后也可能得到新的本地长度。
+- **处理**：让 custom op 边界覆盖产生并消费 data-dependent 中间结果的最小完整语义单元。仅当长度不决定静态 shape、constexpr 控制流或 layout 时，才改为 Triton 运行时标量，并用 `do_not_specialize` 禁止运行时值和对齐关系特化；确需静态的参数使用有限 bucket/padding，不做 request-aware prompt 枚举。
+- **验收**：从冷缓存走完整 `warmup → ready → 正式请求`，确认所有内层编译都发生在 ready 前；用多个长度及非整除 tail 验证数值，且正式 Step 1 回到稳态量级、稳态性能无回退。

.claude/skills/support_model_compile/references/implementation-patterns.md

@@ -13,6 +13,7 @@
 - [Lingbot-Video：normal MoE](#lingbot-videonormal-moe)
 - [Offload 接入](#offload-接入)
 - [第三方算子叶子](#第三方算子叶子)
+- [自定义算子内层编译与 Triton 动态标量](#自定义算子内层编译与-triton-动态标量)
 - [测试骨架](#测试骨架)
 
 ## 参考入口
@@ -183,6 +184,36 @@ return external_kernel(x, weight)
 
 用最小输入验证数值、dtype、shape 和真实 kernel。custom op 不会减少调用次数；需要减少 launch 时应设计更大粒度的融合算子。
 
+边界应覆盖包含 data-dependent 中间结果的最小完整语义单元。当中间结果会迫使 Dynamo 追踪其数据依赖时，block map、增量 LUT 和消费该 LUT 的 attention kernel 应共同留在边界内，不能只把最后一个 kernel 设为叶子；fake 实现只描述该单元对外可见的输出元数据。
+
+## 自定义算子内层编译与 Triton 动态标量
+
+custom op 让外层 Dynamo 只看到一个算子，但函数体仍会正常执行。其内部的 Triton、CUDA extension 和第三方 wrapper 有各自的 JIT 与缓存，因此外层没有 recompile 日志，不代表正式请求没有发生内层编译。服务 ready 前后分别记录内层编译日志或缓存产物，再用冷缓存完整生命周期复测；不要仅凭 `TORCH_LOGS` 判断。
+
+多模态路径按实际执行顺序比较签名：
+
+```text
+原始模态 token
+→ packing / padding / alignment
+→ TP / SP 切分
+→ 各 rank 的实际 q/k 长度
+→ sparse block 数和 leaf kernel 标量参数
+```
+
+prompt 长度可任意变化时，不要靠枚举 prompt 预编译所有长度。先判断 Triton 标量是否必须静态：
+
+- 标量只参与指针偏移、mask 或运行时算术，且 launch grid 可在 JIT 外计算时，可以作为运行时参数；block size 等静态元参数仍保留 `tl.constexpr`。
+- 标量决定 `tl.arange` 范围、静态 tensor shape、constexpr 分支、unroll、layout 或 `num_warps` 时，不能直接动态化；应重新划分参数，或使用明确且有限的 bucket/padding。
+- 对确认安全的长度参数，使用 `do_not_specialize` 禁止按运行时值和对齐关系特化，避免每个长度或整除关系产生新变体：
+
+```python
+@triton.jit(do_not_specialize=("seq_len",))
+def kernel(x, seq_len, block: tl.constexpr):
+    ...
+```
+
+不要批量套用该装饰器。至少用两个不同长度（包含非整除 tail）验证数值，并确认 ready 后不再产生新的内层编译产物，稳态性能也没有回退。
+
 ## 测试骨架
 
 mock `torch.compile` 为原函数，测试 Python cache 语义：

lightx2v/common/ops/attn/dynamic_sparse_attn.py

@@ -30,6 +30,38 @@
     magi_ffa_func = None
 
 
+@torch.library.custom_op(
+    "lightx2v::dynamic_sparse_sage2",
+    mutates_args=(),
+    device_types="cuda",
+)
+def dynamic_sparse_sage2(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    topk_ratio: float,
+    block_q: int,
+    block_k: int,
+    arch: str,
+) -> torch.Tensor:
+    sparse_map, _, _ = get_block_map(q, k, topk_ratio=topk_ratio, BLKQ=block_q, BLKK=block_k)
+    lut, valid_block_num = block_map_incremental_lut_triton(sparse_map)
+    return sage2_block_sparse_attn(q, k, v, lut, valid_block_num, block_q, block_k, arch)
+
+
+@dynamic_sparse_sage2.register_fake
+def _dynamic_sparse_sage2_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    topk_ratio: float,
+    block_q: int,
+    block_k: int,
+    arch: str,
+) -> torch.Tensor:
+    return torch.empty_like(q)
+
+
 @ATTN_WEIGHT_REGISTER("dynamic_sparse_attn")
 class DynamicSparseAttnWeight(AttnWeightTemplate):
     sparsity_ratio = 0.8
@@ -152,10 +184,7 @@ def apply_sage2(
         k = k.unsqueeze(0).transpose(1, 2).contiguous()
         v = v.unsqueeze(0).transpose(1, 2).contiguous()
 
-        sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
-        lut, valid_block_num = block_map_incremental_lut_triton(sparse_map)
-
-        out = sage2_block_sparse_attn(q, k, v, lut, valid_block_num, self.BLKQ, self.BLKK, self.arch)
+        out = dynamic_sparse_sage2(q, k, v, self.topk, self.BLKQ, self.BLKK, self.arch)
         out = out.transpose(1, 2).reshape(max_seqlen_q, -1)
         return out
 

lightx2v/common/ops/attn/utils/sla_util.py

@@ -3,11 +3,11 @@
 import triton.language as tl
 
 
-@triton.jit
+@triton.jit(do_not_specialize=("L",))
 def compress_kernel(
     X,
     XM,
-    L: tl.constexpr,
+    L,
     D: tl.constexpr,
     BLOCK_L: tl.constexpr,
 ):

lightx2v/common/ops/attn/utils/sparge_util.py

@@ -21,7 +21,6 @@
 try:
     import spas_sage_attn._fused as fused
     import spas_sage_attn._qattn as qattn
-    from spas_sage_attn.utils import get_vanilla_qk_quant
 except ImportError:
     print("spas_sage_attn is not installed.")
 
@@ -32,6 +31,63 @@
     SAGE2PP_ENABLED = False
 
 
+@triton.jit(do_not_specialize=("seq_len",))
+def dynamic_qk_quantize_kernel(
+    x_ptr,
+    x_mean_ptr,
+    x_quant_ptr,
+    scale_ptr,
+    seq_len,
+    head_dim: tl.constexpr,
+    block_size: tl.constexpr,
+    subtract_mean: tl.constexpr,
+):
+    batch = tl.program_id(0)
+    head = tl.program_id(1)
+    block = tl.program_id(2)
+    heads = tl.num_programs(1)
+    blocks = tl.num_programs(2)
+
+    row_offsets = block * block_size + tl.arange(0, block_size)
+    dim_offsets = tl.arange(0, head_dim)
+    mask = row_offsets[:, None] < seq_len
+    tensor_offset = (batch * heads + head) * seq_len * head_dim
+    x_ptrs = x_ptr + tensor_offset + row_offsets[:, None] * head_dim + dim_offsets[None, :]
+    x = tl.load(x_ptrs, mask=mask, other=0.0)
+
+    if subtract_mean:
+        mean_ptrs = x_mean_ptr + (batch * heads + head) * head_dim + dim_offsets
+        x = tl.where(mask, x - tl.load(mean_ptrs)[None, :], 0.0)
+
+    x_fp32 = x.to(tl.float32)
+    scale = tl.max(tl.abs(x_fp32)) / 127.0 + 1e-7
+    x_scaled = x_fp32 / scale
+    x_quant = (x_scaled + 0.5 * tl.where(x_scaled >= 0, 1, -1)).to(tl.int8)
+
+    x_quant_ptrs = x_quant_ptr + tensor_offset + row_offsets[:, None] * head_dim + dim_offsets[None, :]
+    tl.store(x_quant_ptrs, x_quant, mask=mask)
+    tl.store(scale_ptr + (batch * heads + head) * blocks + block, scale)
+
+
+def quantize_qk_blocks(x, x_mean, block_size):
+    x = x.contiguous()
+    batch, heads, seq_len, head_dim = x.shape
+    blocks = triton.cdiv(seq_len, block_size)
+    x_quant = torch.empty_like(x, dtype=torch.int8)
+    x_scale = torch.empty((batch, heads, blocks), device=x.device, dtype=torch.float32)
+    dynamic_qk_quantize_kernel[(batch, heads, blocks)](
+        x,
+        x_mean,
+        x_quant,
+        x_scale,
+        seq_len,
+        head_dim,
+        block_size,
+        x_mean is not None,
+    )
+    return x_quant, x_scale
+
+
 def hyperparameter_check(hyper, H, device):
     if type(hyper) is float or type(hyper) is int:
         hyper = torch.full((H,), float(hyper), device=device)
@@ -265,7 +321,8 @@ def sage2_block_sparse_attn(q, k, v, lut, valid_block_num, BLKQ, BLKK, arch):
     assert headdim in [64, 128], "headdim should be in [64, 128]. For other headdim, you can use padding and specify the softmax scale."
 
     km = k.mean(dim=-2, keepdim=True)
-    q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant(q, k, km, BLKQ, BLKK)
+    q_int8, q_scale = quantize_qk_blocks(q, None, BLKQ)
+    k_int8, k_scale = quantize_qk_blocks(k, km, BLKK)
     scale = 1.0 / (headdim**0.5)
 
     o_s = torch.empty_like(q)