[Diffusion] Tune QK head LayerNorm for SM120 (#34349)

python/sglang/kernels/ops/diffusion/triton/layernorm_modulate.py

@@ -46,6 +46,7 @@ import torch
 import triton  # type: ignore
 import triton.language as tl  # type: ignore
 
+from sglang.kernels.jit.utils import get_jit_cuda_arch
 from sglang.kernels.ops.diffusion.triton.numerics import (
     cuda_rsqrtf,
     div_rn_f32,
@@ -340,6 +341,11 @@ def _is_bf16_cuda(t: torch.Tensor) -> bool:
     return t.is_cuda and t.dtype is torch.bfloat16
 
 
+def _is_sm120_or_newer() -> bool:
+    arch = get_jit_cuda_arch()
+    return arch.major * 10 + arch.minor >= 120
+
+
 def _mod_row_stride(t: torch.Tensor, batch: int, hidden: int) -> int | None:
     # (batch, hidden) modulation rows, possibly strided views of a chunked
     # adaLN projection; the last dim must be packed.
@@ -454,7 +460,12 @@ def fused_qk_head_layernorm(
     launch, bit-exact vs the eager aten kernel."""
     head_dim = q.shape[-1]
     n_rows = q.numel() // head_dim
-    rows = 64
+    # SM120 has a smaller register file per SM than H200.  Grouping 64 exact
+    # Welford rows in one program depresses occupancy on RTX 5090; an exhaustive
+    # production-shape sweep selects 8 rows / 4 warps there (about 10% faster).
+    # Preserve the independently tuned SM90 launch byte-for-byte.
+    is_sm120 = _is_sm120_or_newer()
+    rows = 8 if is_sm120 else 64
     q_out = torch.empty_like(q)
     k_out = torch.empty_like(k)
     with torch.cuda.device(q.device):
@@ -470,6 +481,6 @@ def fused_qk_head_layernorm(
             ROWS=rows,
             # H200-tuned: 62us at (1, 4360, 32, 128) vs the 301us of the two
             # aten launches (one 128-thread block per head_dim-element row).
-            num_warps=2,
+            num_warps=4 if is_sm120 else 2,
         )
     return q_out, k_out