chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[jit kernel] support dtype as a cpp template parameter (#16452)

Browse Source
This commit is contained in:
陈一涵
2026-01-08 13:54:33 +08:00
committed by GitHub
parent 41b434a7e6
commit 48b8dcd42e
9 changed files with 91 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

@@ -445,7 +445,7 @@ def apply_qk_norm(
q.is_cuda
and allow_inplace
and (q_eps == k_eps)
and can_use_fused_inplace_qknorm(head_dim)
and can_use_fused_inplace_qknorm(head_dim, q.dtype)
):
fused_inplace_qknorm(
q=q.view(batch_size, -1, head_dim),

2
python/sglang/multimodal_gen/runtime/models/dits/zimage.py
View File

@@ -159,7 +159,7 @@ class ZImageAttention(nn.Module):
if (
q.is_cuda
and (self.norm_q.variance_epsilon == self.norm_k.variance_epsilon)
and can_use_fused_inplace_qknorm(self.head_dim)
and can_use_fused_inplace_qknorm(self.head_dim, q.dtype)
):
q, k = apply_qk_norm(
q=q,