diff --git a/.claude/skills/support_model_compile/SKILL.md b/.claude/skills/support_model_compile/SKILL.md index e006584fe..ef0b30214 100644 --- a/.claude/skills/support_model_compile/SKILL.md +++ b/.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 + 正式请求的功能验收； - 每组至少三轮； diff --git a/.claude/skills/support_model_compile/references/casebook.md b/.claude/skills/support_model_compile/references/casebook.md index a2823351c..addfa5f5c 100644 --- a/.claude/skills/support_model_compile/references/casebook.md +++ b/.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 回到稳态量级、稳态性能无回退。 diff --git a/.claude/skills/support_model_compile/references/implementation-patterns.md b/.claude/skills/support_model_compile/references/implementation-patterns.md index 26ad0a1bf..24235ce47 100644 --- a/.claude/skills/support_model_compile/references/implementation-patterns.md +++ b/.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 语义： diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index ff720f0f2..eb9281e81 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/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 diff --git a/lightx2v/common/ops/attn/utils/sla_util.py b/lightx2v/common/ops/attn/utils/sla_util.py index d44c77556..d717d9a94 100755 --- a/lightx2v/common/ops/attn/utils/sla_util.py +++ b/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, ): diff --git a/lightx2v/common/ops/attn/utils/sparge_util.py b/lightx2v/common/ops/attn/utils/sparge_util.py index 9aa2e777b..bd94cbed4 100644 --- a/lightx2v/common/ops/attn/utils/sparge_util.py +++ b/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)