Skip to content

Optimize Gemma4 H200 MoE and extend attention #26588

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
BBuf wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Optimize Gemma4 H200 MoE and extend attention #26588

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
3398cb7
Optimize Gemma4 H200 MoE and extend attention
BBuf May 27, 2026
03826cd
Optimize Gemma4 MoE routing topk
BBuf May 27, 2026
d72d246
Fuse Gemma4 router RMSNorm scale
BBuf May 27, 2026
fabb5b5
Optimize Gemma4 QKV RMSNorm for small batches
BBuf May 27, 2026
256e1d6
Merge remote-tracking branch 'bbuf/main' into codex/gemma4-h200-moe-attn
BBuf May 28, 2026
67f9d13
Merge remote-tracking branch 'origin/main' into codex/gemma4-h200-moe…
May 30, 2026
10ab189
Merge remote-tracking branch 'origin/main' into codex/gemma4-h200-moe…
May 30, 2026
0d3c9c6
Revert "Fuse Gemma4 router RMSNorm scale"
May 31, 2026
27cb94c
Revert "Optimize Gemma4 MoE routing topk"
May 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion python/sglang/srt/layers/attention/triton_ops/extend_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,10 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
BLOCK_M, BLOCK_N = (16, 64)
elif _is_cuda and CUDA_CAPABILITY[0] >= 9:
# Hopper architecture (H100, etc.)
if Lq <= 256:
if Lq <= 128:
BLOCK_M, BLOCK_N = (128, 64)
elif Lq <= 256:
BLOCK_M, BLOCK_N = (64, 64)
else:
BLOCK_M, BLOCK_N = (32, 64)
elif _is_cuda and CUDA_CAPABILITY[0] >= 8:
Expand Down
66 changes: 66 additions & 0 deletions python/sglang/srt/layers/gemma4_fused_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,52 @@ def _gemma_qkv_rmsnorm_kernel(
tl.store(V_ptr + off, out.to(V_ptr.dtype.element_ty), mask=mask)


@triton.jit
def _gemma_qkv_rmsnorm_by_head_kernel(
Q_ptr,
K_ptr,
V_ptr,
Q_w_ptr,
K_w_ptr,
stride_q_m,
stride_k_m,
stride_v_m,
NUM_Q_HEADS: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
HEAD_DIM: tl.constexpr,
eps,
HAS_KV: tl.constexpr,
BLOCK: tl.constexpr,
):
m = tl.program_id(0)
h_all = tl.program_id(1)
cols = tl.arange(0, BLOCK)
mask = cols < HEAD_DIM

if h_all < NUM_Q_HEADS:
off = m * stride_q_m + h_all * HEAD_DIM + cols
x = tl.load(Q_ptr + off, mask=mask, other=0.0).to(tl.float32)
qw = tl.load(Q_w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps)
out = x * rrms * qw
tl.store(Q_ptr + off, out.to(Q_ptr.dtype.element_ty), mask=mask)
elif HAS_KV and h_all < NUM_Q_HEADS + NUM_KV_HEADS:
h = h_all - NUM_Q_HEADS
off = m * stride_k_m + h * HEAD_DIM + cols
x = tl.load(K_ptr + off, mask=mask, other=0.0).to(tl.float32)
kw = tl.load(K_w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps)
out = x * rrms * kw
tl.store(K_ptr + off, out.to(K_ptr.dtype.element_ty), mask=mask)
elif HAS_KV:
h = h_all - NUM_Q_HEADS - NUM_KV_HEADS
off = m * stride_v_m + h * HEAD_DIM + cols
x = tl.load(V_ptr + off, mask=mask, other=0.0).to(tl.float32)
rrms = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps)
out = x * rrms
tl.store(V_ptr + off, out.to(V_ptr.dtype.element_ty), mask=mask)


def gemma_qkv_rmsnorm(
q: torch.Tensor,
k: Optional[torch.Tensor],
Expand Down Expand Up @@ -227,6 +273,26 @@ def gemma_qkv_rmsnorm(
assert k.stride(-1) == 1 and v.stride(-1) == 1
assert k_weight is not None and k_weight.shape[-1] == head_dim

if M <= 256:
total_heads = num_q_heads + (2 * num_kv_heads if has_kv else 0)
_gemma_qkv_rmsnorm_by_head_kernel[(M, total_heads)](
q,
k if has_kv else q,
v if has_kv else q,
q_weight,
k_weight if has_kv else q_weight,
q.stride(0),
k.stride(0) if has_kv else 0,
v.stride(0) if has_kv else 0,
NUM_Q_HEADS=num_q_heads,
NUM_KV_HEADS=num_kv_heads if has_kv else 0,
HEAD_DIM=head_dim,
eps=eps,
HAS_KV=has_kv,
BLOCK=BLOCK,
)
return

_gemma_qkv_rmsnorm_kernel[(M,)](
q,
k if has_kv else q,
Expand Down
114 changes: 114 additions & 0 deletions ...moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=704,device_name=NVIDIA_H200.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2
},
"512": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
}
}
114 changes: 114 additions & 0 deletions ...oe_runner/triton_utils/configs/triton_3_6_0/E=128,N=704,device_name=NVIDIA_H200_down.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2
},
"512": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
}
}
Loading