From 26579d893e4a7405fcee43479cd0afbe8291d0af Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 13 Aug 2026 10:49:13 +0800 Subject: [PATCH] [Diffusion][GLM-Image] Retune QK head LayerNorm for SM103 (#34619) --- .../sglang/kernels/ops/diffusion/triton/layernorm_modulate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py index ead6d347e..91314b519 100644 --- a/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py +++ b/python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py @@ -344,7 +344,7 @@ def _is_bf16_cuda(t: torch.Tensor) -> bool: def _qk_head_launch_config() -> tuple[int, int]: arch = get_jit_cuda_arch() if arch.major == 10 and arch.minor == 3: - return 16, 1 + return 32, 1 if arch.major * 10 + arch.minor >= 120: return 8, 4 return 64, 2 @@ -464,7 +464,7 @@ def fused_qk_head_layernorm( launch, bit-exact vs the eager aten kernel.""" head_dim = q.shape[-1] n_rows = q.numel() // head_dim - # Architecture sweeps at the production GLM shape select 16 rows / 1 warp + # Architecture sweeps at the production GLM shape select 32 rows / 1 warp # on B300 (SM103) and 8 rows / 4 warps on RTX 5090 (SM120). Preserve the # independently tuned H100/H200 launch on all other architectures. rows, num_warps = _qk_head_launch_config()