minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Add Work-Centric (Lean) Attention: a persistent-CTA decode kernel for long-context serving (#33576)

Browse Source 
Signed-off-by: Valerie Chen <valechen@amd.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
valechen
2026-08-28 17:00:44 -07:00
committed by GitHub
co-authored by HAI
parent db6f0a9d53
commit 4944e50e2c
6 changed files with 1585 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/lean_kernel_sweep.py
Executable
+163
View File
@@ -0,0 +1,163 @@
#!/usr/bin/env python3
"""
Kernel-level Lean vs Standard (SplitK) decode-attention sweep, both GQA head
configs, full batch x context grid, at the shipped 1xCU persistent grid.
Generalizes benchmark/lean_kernel_qwen_gqa.py (which is Qwen batch=1 only) to
sweep Qwen2.5-7B (28Q/4KV) and Llama-3.1-8B (32Q/8KV) over
batch in {1,2,4,8,16,32} x context in {8K,16K,32K,64K,128K}, reporting per-call
kernel latency, speedup (std / lean), cosine parity vs SplitK, and whether the
eager auto-gate would enable Lean. Writes a CSV for the PR tables.
Grid is whatever SGLANG_FORCE_LEAN_GRID_CU_MULT resolves to (default 1.0 = one
CTA per CU); set it to A/B other grids without a rebuild.
Usage: python3 benchmark/lean_kernel_sweep.py [out.csv]
"""
import os
import sys
import time
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
import torch
from sglang.kernels.ops.attention.decode_attention import (
_LEAN_BLOCK_M,
_lean_decode_launch_params,
decode_attention_fwd,
decode_attention_fwd_grouped,
lean_decode_seqlen_gate,
)
MODELS = [
("qwen2.5-7b", 28, 4),
("llama3.1-8b", 32, 8),
]
D = D_V = 128
MAX_KV_SPLITS = 8
BATCHES = [1, 2, 4, 8, 16, 32]
CONTEXTS = [8192, 16384, 32768, 65536, 131072]
def bench(fn, warmup=15, iters=100):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(iters):
fn()
torch.cuda.synchronize()
return (time.perf_counter() - t0) / iters * 1000.0 # ms
def run(H_Q, H_KV, B, S):
dev, dt = "cuda", torch.float16
kvg = H_Q // H_KV
sm = 1.0 / (D**0.5)
tot = B * S
total_programs, _, _ = _lean_decode_launch_params(H_KV, kvg)
lean_Mp = torch.empty(
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
)
lean_Lp = torch.empty(
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
)
lean_Op = torch.empty(
(total_programs, _LEAN_BLOCK_M, D_V), dtype=torch.float32, device=dev
)
lean_locks = torch.zeros((total_programs,), dtype=torch.int32, device=dev)
kvi = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
kvx = torch.arange(0, tot, device=dev, dtype=torch.int32)
q = torch.randn(B, H_Q, D, dtype=dt, device=dev)
k = torch.randn(tot, H_KV, D, dtype=dt, device=dev)
v = torch.randn(tot, H_KV, D_V, dtype=dt, device=dev)
attn_logits = torch.empty(
(B, H_Q, MAX_KV_SPLITS, D_V), dtype=torch.float32, device=dev
)
attn_lse = torch.empty((B, H_Q, MAX_KV_SPLITS), dtype=torch.float32, device=dev)
nks = torch.full((B,), MAX_KV_SPLITS, dtype=torch.int32, device=dev)
o_std = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
std_ms = bench(
lambda: decode_attention_fwd_grouped(
q, k, v, o_std, kvi, kvx, attn_logits, attn_lse, nks, MAX_KV_SPLITS, sm, 1.0
)
)
attn_logits2 = torch.empty_like(attn_logits)
attn_lse2 = torch.empty_like(attn_lse)
o_lean = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
lean_ms = bench(
lambda: decode_attention_fwd(
q,
k,
v,
o_lean,
kvi,
kvx,
attn_logits2,
attn_lse2,
nks,
MAX_KV_SPLITS,
sm,
1.0,
1.0,
enable_lean=True,
lean_Mp=lean_Mp,
lean_Lp=lean_Lp,
lean_Op=lean_Op,
lean_locks=lean_locks,
)
)
cos = torch.nn.functional.cosine_similarity(
o_lean.flatten().float(), o_std.flatten().float(), dim=0
).item()
gate = lean_decode_seqlen_gate(H_Q, kvg, B, B * S, is_mla=False)
return std_ms, lean_ms, cos, gate
def main():
out = sys.argv[1] if len(sys.argv) > 1 else "grid_out/kernel_sweep_1xcu.csv"
mult = float(os.environ.get("SGLANG_FORCE_LEAN_GRID_CU_MULT", "1.0"))
print(f"\nKernel sweep GPU={torch.cuda.get_device_name(0)} grid_mult={mult}")
print("=" * 78)
rows = ["model,H_Q,H_KV,batch,context,std_ms,lean_ms,speedup,cos,gate"]
for name, H_Q, H_KV in MODELS:
print(f"\n{name} ({H_Q}Q/{H_KV}KV)")
print(
f"{'batch':>5} {'ctx':>6} {'std_ms':>9} {'lean_ms':>9} {'speedup':>8} {'gate':>5} {'cos':>7}"
)
for B in BATCHES:
for S in CONTEXTS:
# b32 x 128K on the 8-KV-head config exceeds the microbench's single
# contiguous KV tensor (faults the GPU); real serving uses a paged pool.
if H_KV == 8 and B == 32 and S == 131072:
print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}")
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
continue
try:
std, lean, cos, gate = run(H_Q, H_KV, B, S)
except torch.cuda.OutOfMemoryError:
torch.cuda.empty_cache()
print(f"{B:>5} {S//1024:>5}K {'OOM':>9}")
rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,")
continue
sp = std / lean
print(
f"{B:>5} {S//1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
)
rows.append(
f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}"
)
os.makedirs(os.path.dirname(out), exist_ok=True)
open(out, "w").write("\n".join(rows) + "\n")
print(f"\nwrote {out}")
if __name__ == "__main__":
main()
python/sglang/kernels/ops/attention/decode_attention.py
+832
View File
@@ -21,8 +21,10 @@ It supports page size = 1.
# https://github.com/ModelTC/lightllm/blob/96353e868a840db4d103138caf15ed9dbea8c186/lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage2.py
import logging
import math
from typing import NamedTuple, Optional, Tuple
import torch
import triton
import triton.language as tl
@@ -1182,6 +1184,11 @@ def decode_attention_fwd(
page_size: int = 1,
score_mod=None,
aux_tensors=None,
enable_lean=None,
lean_Mp=None,
lean_Lp=None,
lean_Op=None,
lean_locks=None,
):
assert max_kv_splits == attn_logits.shape[2]
assert q.shape[0] <= kv_indptr.shape[0] - 1
@@ -1191,6 +1198,47 @@ def decode_attention_fwd(
kv_head_num = v_buffer.shape[-2]
kv_group_num = q.shape[1] // kv_head_num
# Work-Centric (Lean) Attention: a persistent-CTA + work-stealing decode kernel
# that helps on long sequences where there are many more KV tiles than CUs. The
# persistent grid is fixed to the device CU count and the kernel derives its own tile
# schedule from kv_indptr on-device, so this path involves no host sync and is safe to
# capture in a CUDA graph. Whether Lean pays off for a given shape is decided cheaply by
# the backend's host-side seqlen gate (lean_decode_seqlen_gate) before we get here.
# Lean supports both the contiguous 3-D [N, head, dim] and paged 4-D
# [num_pages, page_size, head, dim] KV layouts (page-aware address math in the kernel).
# ROCm/AMD only: Lean is validated on MI300X/MI355X; CUDA/NVIDIA uses the standard kernel.
if (
_is_hip
and _lean_head_dim_ok(k_buffer.shape[-1], v_buffer.shape[-1])
and _should_use_lean_decode(
enable_lean, logit_cap, sinks, xai_temperature_len, score_mod
)
):
total_programs, XCD_REMAP, NUM_XCDS = _lean_decode_launch_params(
v_buffer.shape[-2], kv_group_num
)
_decode_lean_attention_fwd(
q,
k_buffer,
v_buffer,
o,
kv_indptr,
kv_indices,
total_programs,
# Fold k_scale into sm_scale and pass v_scale, exactly as the standard grouped
# kernel does, so Lean dequantizes fp8 KV consistently (both are 1.0 for bf16/fp16).
sm_scale * k_scale,
v_scale,
XCD_REMAP,
NUM_XCDS,
lean_Mp,
lean_Lp,
lean_Op,
lean_locks,
page_size=page_size,
)
return
if kv_group_num == 1:
# MHA
decode_attention_fwd_normal(
@@ -1237,3 +1285,787 @@ def decode_attention_fwd(
score_mod=score_mod,
aux_tensors=aux_tensors,
)
# ============================================================================
# Work-Centric (Lean) Attention: persistent-CTA + work-stealing decode kernel.
# ============================================================================
_LEAN_BLOCK_M = 16
_NUM_CU = None
def _lean_head_dim_ok(qk_head_dim: int, v_head_dim: int) -> bool:
"""Whether the Lean decode kernel's tiles fit in shared memory for this head dim.
The non-MLA kernel sets ``BLOCK_DMODEL = next_power_of_2(qk_head_dim)``; at head_dim 256
(e.g. Gemma-2/3) the K/V tiles overflow the 160 KB LDS budget and the launch raises
OutOfResources. head_dim <= 128 fits. MLA's rope-split dims (288/576) are special-cased in
the kernel into a smaller-tiled path and are handled separately. This guard makes the Lean
dispatch fall back safely instead of crashing, even under an explicit ``enable_lean=True``.
"""
if qk_head_dim in (288, 576): # MLA rope-split, special-cased in the kernel
return True
return qk_head_dim <= 128 and v_head_dim <= 128
def _lean_num_cus() -> int:
"""Number of compute units on the current device (cached).
Lean Attention sizes its persistent grid to the hardware CU count so work-stealing can
fill the GPU. Falls back to 304 (MI300X) if the device cannot be queried.
"""
global _NUM_CU
if _NUM_CU is None:
try:
_NUM_CU = torch.cuda.get_device_properties(0).multi_processor_count
except Exception:
_NUM_CU = 304
return _NUM_CU
def _lean_decode_block_n(Lk: int) -> int:
"""KV block size for the Lean decode kernel.
Large head dims (MLA, Lk in {288, 576}) use a small KV block to bound LDS/register
usage; standard head dims use a large block since decode is memory-bound. The value
must be identical everywhere it is used so the tile schedule stays consistent.
"""
if not _is_hip:
return 64
return 16 if Lk > 256 else 128
@triton.jit
def remap_xcd(pid, GRID_MN: tl.constexpr, NUM_XCDS: tl.constexpr = 8):
"""Remap program ID across XCDs for AMD MI300X."""
pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
tall_xcds = GRID_MN % NUM_XCDS
tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
xcd = pid % NUM_XCDS
local_pid = pid // NUM_XCDS
if xcd < tall_xcds:
pid = xcd * pids_per_xcd + local_pid
else:
pid = (
tall_xcds * pids_per_xcd
+ (xcd - tall_xcds) * (pids_per_xcd - 1)
+ local_pid
)
return pid, pids_per_xcd
@triton.jit
def cal_num_split_wgs(
xcd_pid: tl.int32,
tile_iter_end: tl.int32,
cta_end_tile_gid: tl.int32,
max_tiles_per_wg: tl.int32,
high_load_wgs: tl.int32,
num_splits: tl.int32,
):
zero_i = tl.full((), 0, dtype=tl.int32)
start_cta = tl.cast(xcd_pid + 1, tl.int32)
remaining = tl.maximum(tl.cast(tile_iter_end - cta_end_tile_gid, tl.int32), zero_i)
cap_high = tl.cast(max_tiles_per_wg, tl.int32)
cap_low = tl.cast(max_tiles_per_wg - 1, tl.int32)
cap_low = tl.where(cap_low > 0, cap_low, tl.full((), 1, dtype=tl.int32))
ctas_high_avail = tl.maximum(tl.cast(high_load_wgs, tl.int32) - start_cta, zero_i)
total_high_capacity = ctas_high_avail * cap_high
need_high_only = (remaining + cap_high - 1) // cap_high
rem_after_high = tl.maximum(remaining - total_high_capacity, zero_i)
need_low_after_high = (rem_after_high + cap_low - 1) // cap_low
ctas_needed = tl.where(
remaining <= total_high_capacity,
need_high_only,
ctas_high_avail + need_low_after_high,
)
max_ctas_allowed = tl.maximum(tl.cast(num_splits - 1, tl.int32), zero_i)
ctas_to_use = tl.minimum(ctas_needed, max_ctas_allowed)
k = ctas_to_use
cap_by_k = tl.where(
k <= ctas_high_avail,
k * cap_high,
total_high_capacity + (k - ctas_high_avail) * cap_low,
)
last_cta = start_cta + ctas_to_use
last_cta = tl.where(ctas_to_use == 0, start_cta - 1, last_cta)
return last_cta
@triton.jit
def _lean_attention_decode_kernel(
Q,
K_Buffer,
V_Buffer,
Mp, # Partial max
Lp, # Partial sum
Op, # Partial output
O, # Final output
batch_num_block_n,
locks,
kv_indptr,
kv_indices,
sm_scale,
v_scale,
stride_qbs,
stride_qh,
stride_buf_kbs,
stride_buf_kh,
stride_buf_kpage,
stride_buf_ktok,
stride_buf_vbs,
stride_buf_vh,
stride_buf_vpage,
stride_buf_vtok,
stride_obs,
stride_oh,
kv_group_num: tl.constexpr,
NUM_HEAD_BLOCKS: tl.constexpr,
ROWS_PER_XCD: tl.constexpr,
BLOCK_DMODEL: tl.constexpr,
BLOCK_DPE: tl.constexpr,
BLOCK_DV: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
PAGE_SIZE: tl.constexpr,
XCD_REMAP: tl.constexpr,
NUM_XCDS: tl.constexpr,
batch_size: tl.constexpr,
total_programs: tl.constexpr,
num_query_heads: tl.constexpr,
num_rows: tl.constexpr,
xcd_programs: tl.constexpr,
max_output_tile_cnt: tl.constexpr,
Lk: tl.constexpr,
Lv: tl.constexpr,
):
"""Lean Attention decode kernel - persistent CTA with work stealing.
The tile schedule (``tiles_per_khead``, ``max_tiles_per_wg``, ``high_load_wgs``,
``num_splits``) is computed here on-device from ``kv_indptr`` rather than passed in from
the host. This keeps the launch free of any host sync (so it is CUDA-graph capturable)
and lets the schedule adapt to the per-replay sequence length: ``total_programs`` is a
fixed persistent grid and the work simply re-distributes when the KV length changes.
"""
current_pid = tl.program_id(0)
# On-device tile schedule (mirrors the former host-side la_get_num_splits). Reads only
# GPU state so it is safe under CUDA-graph capture. tiles_per_khead is the number of KV
# tiles in one row summed over the batch: it MUST match batch_num_block_n (the per-batch
# cumulative tile count) exactly, so it is read from that array's last entry rather than
# recomputed as ceil(total_tokens / BLOCK_N) -- those differ whenever a sequence length
# is not a multiple of BLOCK_N (the common case for a ragged decode batch), which would
# desync the row<->tile mapping below.
tiles_per_khead = tl.load(batch_num_block_n + batch_size - 1)
# Effective rows per XCD (constexpr-folded); total tiles distributed over this XCD.
eff_rows: tl.constexpr = num_rows // NUM_XCDS if XCD_REMAP else num_rows
total_tiles = tiles_per_khead * eff_rows
max_tiles_per_wg = (total_tiles + xcd_programs - 1) // xcd_programs
max_tiles_per_wg = tl.maximum(max_tiles_per_wg, 1)
high_load_wgs = total_tiles - (max_tiles_per_wg - 1) * xcd_programs
# Safe over-estimate of the split count: a row spans at most ceil(tiles/(mtpw-1))+1
# CTAs; the guarded divisor also covers the max_tiles_per_wg == 1 case.
split_denom = tl.maximum(max_tiles_per_wg - 1, 1)
num_splits = 1 + (tiles_per_khead + split_denom - 1) // split_denom
if XCD_REMAP:
current_pid, pids_per_xcd = remap_xcd(
current_pid, GRID_MN=total_programs, NUM_XCDS=NUM_XCDS
)
xcd_pid = current_pid % pids_per_xcd
xcd_id = current_pid // pids_per_xcd
else:
xcd_pid = current_pid
xcd_id = 0
pids_per_xcd = total_programs
if xcd_pid < high_load_wgs:
iter = max_tiles_per_wg * xcd_pid
cta_end_tile_gid = iter + max_tiles_per_wg
else:
iter = (max_tiles_per_wg - 1) * (
xcd_pid - high_load_wgs
) + high_load_wgs * max_tiles_per_wg
cta_end_tile_gid = iter + (max_tiles_per_wg - 1)
# Use a regular while loop instead of tl.static_range with a dynamic bound to avoid
# Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed).
while iter < cta_end_tile_gid:
tile_row_idx = iter // tiles_per_khead
tile_idx = tile_row_idx * batch_size
tile_iter = tile_row_idx * tiles_per_khead
if batch_size == 1:
req_size = tl.full((), tiles_per_khead, dtype=tl.int32)
else:
req_size = tl.cast(tl.load(batch_num_block_n), tl.int32)
tile_iter_end = tile_iter + req_size
for b in range(1, batch_size):
next_req_size = tl.load(batch_num_block_n + b)
local_head_iter = iter % tiles_per_khead
if (local_head_iter < next_req_size) and (local_head_iter >= req_size):
tile_iter = tile_iter + req_size
tile_idx = tile_idx + b
tile_iter_end = tile_iter + (next_req_size - req_size)
req_size = next_req_size
local_iter = iter - tile_iter
local_iter_end = tl.minimum(tile_iter_end, cta_end_tile_gid) - tile_iter
host_block = iter == tile_iter
finishing_block = cta_end_tile_gid >= tile_iter_end
# A tiling "row" is a (kv_head, head_block) pair. For MHA/GQA NUM_HEAD_BLOCKS == 1
# so a row is just a kv head. For MLA, kv_group_num > BLOCK_M, so a kv head spans
# NUM_HEAD_BLOCKS head blocks of BLOCK_M query heads each.
tile_row_idx_global = ROWS_PER_XCD * xcd_id + tile_row_idx
cur_kv_head = tile_row_idx_global // NUM_HEAD_BLOCKS
head_block_idx = tile_row_idx_global % NUM_HEAD_BLOCKS
group_start = cur_kv_head * kv_group_num
q_head_base = group_start + head_block_idx * BLOCK_M
tile_batch_idx = tile_idx % batch_size
cur_batch = tile_batch_idx
cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
# SGLang-style offsets
offs_d = tl.arange(0, BLOCK_DMODEL)
offs_dv = tl.arange(0, BLOCK_DV)
mask_d = offs_d < Lk
mask_dv = offs_dv < Lv
# Query head block: this row covers BLOCK_M query heads of its kv group, bounded
# by the group end (group_start + kv_group_num) and the total head count.
offs_h = q_head_base + tl.arange(0, BLOCK_M)
mask_h = offs_h < (group_start + kv_group_num)
mask_h = mask_h & (offs_h < num_query_heads)
off_q = cur_batch * stride_qbs + offs_h[:, None] * stride_qh + offs_d[None, :]
q = tl.load(
Q + off_q, mask=mask_h[:, None] & mask_d[None, :], other=0.0
) # [BLOCK_M, BLOCK_DMODEL]
# Cast q to the K buffer dtype so the main dot is a same-dtype MMA. For fp8 KV this
# makes it dot(fp8, fp8) (triton rejects a bf16xfp8 mix); k_scale is folded into
# sm_scale to dequantize. For bf16/fp16 KV this is a no-op. Mirrors the standard kernel.
q_k = q.to(K_Buffer.dtype.element_ty)
# MLA rope split: the positional-encoding dims live in [BLOCK_DMODEL, Lk).
if BLOCK_DPE > 0:
offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
mask_dpe = offs_dpe < Lk
off_qpe = (
cur_batch * stride_qbs + offs_h[:, None] * stride_qh + offs_dpe[None, :]
)
qpe = tl.load(
Q + off_qpe, mask=mask_h[:, None] & mask_dpe[None, :], other=0.0
)
e_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
e_sum = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_DV], dtype=tl.float32)
local_iter_ptr = local_iter * BLOCK_N
local_iter_end_ptr = local_iter_end * BLOCK_N
# Effective token bound: the last tile of a sequence whose length is not a multiple
# of BLOCK_N is only partially valid. Clamp to cur_batch_seq_len so the KV-index /
# K / V loads never read past this batch's tokens -- for the final batch that would
# otherwise run off the end of kv_indices and fault the GPU. For BLOCK_N-aligned
# sequences this equals local_iter_end_ptr, so the aligned path is unchanged.
tok_end = tl.minimum(local_iter_end_ptr, cur_batch_seq_len)
for start_n in range(local_iter_ptr, local_iter_end_ptr, BLOCK_N):
offs_n = start_n + tl.arange(0, BLOCK_N)
kv_loc = tl.load(
kv_indices + cur_batch_kv_start_idx + offs_n,
mask=offs_n < tok_end,
other=0,
)
# Load K transposed: [BLOCK_DMODEL, BLOCK_N] so qk = q @ k directly.
# Page-aware KV address math (mirrors the standard grouped kernel): at
# PAGE_SIZE==1 the slot index addresses directly; otherwise it splits into
# (page_id, tok_in_p) for a [num_pages, page_size, head, dim] paged buffer.
if PAGE_SIZE == 1:
offs_buf_k = (
kv_loc[None, :] * stride_buf_kbs
+ cur_kv_head * stride_buf_kh
+ offs_d[:, None]
)
else:
page_id = kv_loc // PAGE_SIZE
tok_in_p = kv_loc % PAGE_SIZE
offs_buf_k = (
page_id[None, :] * stride_buf_kpage
+ tok_in_p[None, :] * stride_buf_ktok
+ cur_kv_head * stride_buf_kh
+ offs_d[:, None]
)
k = tl.load(
K_Buffer + offs_buf_k,
mask=(offs_n[None, :] < tok_end) & (mask_d[:, None]),
other=0.0,
)
qk = tl.dot(q_k, k) # [BLOCK_M, BLOCK_N]
if BLOCK_DPE > 0:
if PAGE_SIZE == 1:
offs_buf_kpe = (
kv_loc[None, :] * stride_buf_kbs
+ cur_kv_head * stride_buf_kh
+ offs_dpe[:, None]
)
else:
offs_buf_kpe = (
page_id[None, :] * stride_buf_kpage
+ tok_in_p[None, :] * stride_buf_ktok
+ cur_kv_head * stride_buf_kh
+ offs_dpe[:, None]
)
kpe = tl.load(
K_Buffer + offs_buf_kpe,
mask=(offs_n[None, :] < tok_end) & (mask_dpe[:, None]),
other=0.0,
)
# Dequantize the rope-split K to q's dtype for this small dot (matches standard).
qk += tl.dot(qpe, kpe.to(qpe.dtype))
qk *= sm_scale # sm_scale carries k_scale (folded by the caller)
qk = tl.where(
mask_h[:, None] & (offs_n[None, :] < tok_end),
qk,
float("-inf"),
)
n_e_max = tl.maximum(tl.max(qk, 1), e_max)
re_scale = tl.exp(e_max - n_e_max)
p = tl.exp(qk - n_e_max[:, None])
if PAGE_SIZE == 1:
offs_buf_v = (
kv_loc[:, None] * stride_buf_vbs
+ cur_kv_head * stride_buf_vh
+ offs_dv[None, :]
)
else:
offs_buf_v = (
page_id[:, None] * stride_buf_vpage
+ tok_in_p[:, None] * stride_buf_vtok
+ cur_kv_head * stride_buf_vh
+ offs_dv[None, :]
)
v = tl.load(
V_Buffer + offs_buf_v,
mask=(offs_n[:, None] < tok_end) & (mask_dv[None, :]),
other=0.0,
)
acc *= re_scale[:, None]
acc += tl.dot(p.to(v.dtype), v) # [BLOCK_M, BLOCK_DV]
e_sum = e_sum * re_scale + tl.sum(p, 1)
e_max = n_e_max
if not host_block:
mp_ptrs = Mp + current_pid * BLOCK_M + tl.arange(0, BLOCK_M)
lp_ptrs = Lp + current_pid * BLOCK_M + tl.arange(0, BLOCK_M)
op_ptrs = (
Op
+ current_pid * BLOCK_M * BLOCK_DV
+ tl.arange(0, BLOCK_M)[:, None] * BLOCK_DV
+ offs_dv[None, :]
)
tl.store(mp_ptrs, e_max, cache_modifier=".wb")
tl.store(lp_ptrs, e_sum, cache_modifier=".wb")
tl.store(op_ptrs, acc, mask=mask_dv[None, :], cache_modifier=".wb")
tl.debug_barrier()
tl.atomic_xchg(locks + current_pid, 1)
else:
if not finishing_block:
last_cta = cal_num_split_wgs(
xcd_pid=xcd_pid,
tile_iter_end=tile_iter_end,
cta_end_tile_gid=cta_end_tile_gid,
max_tiles_per_wg=max_tiles_per_wg,
high_load_wgs=high_load_wgs,
num_splits=num_splits,
)
# Defensive clamp: the partial-result buffers (Mp/Lp/Op/locks) hold one slot
# per program, and a CTA only ever steals from later CTAs within its own XCD.
# Clamp to pids_per_xcd so a degenerate schedule (e.g. a forced tiny shape
# that slips past the host gate) can never index a buffer out of bounds.
last_cta = tl.minimum(last_cta, pids_per_xcd)
temp_pid = current_pid
for cta in range((xcd_pid + 1), last_cta):
temp_pid = temp_pid + 1
while tl.atomic_cas(locks + temp_pid, 1, 1) != 1:
pass
mp_ptrs = Mp + temp_pid * BLOCK_M + tl.arange(0, BLOCK_M)
lp_ptrs = Lp + temp_pid * BLOCK_M + tl.arange(0, BLOCK_M)
op_ptrs = (
Op
+ temp_pid * BLOCK_M * BLOCK_DV
+ tl.arange(0, BLOCK_M)[:, None] * BLOCK_DV
+ offs_dv[None, :]
)
m_cta = tl.load(mp_ptrs)
l_cta = tl.load(lp_ptrs)
acc_cta = tl.load(op_ptrs, mask=mask_dv[None, :])
m_new = tl.maximum(m_cta, e_max)
alpha = tl.exp(m_cta - m_new)
alpha1 = tl.exp(e_max - m_new)
l_new = alpha * l_cta + alpha1 * e_sum
acc = acc_cta * alpha[:, None] + acc * alpha1[:, None]
e_max = m_new
e_sum = l_new
# v_scale dequantizes the fp8 V contribution accumulated via dot(p, v); it is 1.0
# for bf16/fp16 V. Applied once here at the single output-write site (mirrors the
# standard kernel's `acc / e_sum * v_scale`).
acc = acc / e_sum[:, None] * v_scale
offs_o = (
cur_batch * stride_obs + offs_h[:, None] * stride_oh + offs_dv[None, :]
)
tl.store(O + offs_o, acc, mask=mask_h[:, None] & mask_dv[None, :])
iter = iter + (local_iter_end - local_iter)
def _lean_head_tiles(num_q_heads: int, kv_group_num: int) -> int:
"""Head-tile programs the standard grouped decode kernel launches per (sequence,
kv-split): ``ceil(num_q_heads / min(16, kv_group_num))``. This is the standard
kernel's query-head parallelism, which drives how well it already fills the device
and hence where the Lean-vs-SplitK crossover sits (see :func:`lean_decode_seqlen_gate`).
"""
block_h = min(16, max(1, kv_group_num))
return -(-num_q_heads // block_h) # ceil(num_q_heads / block_h)
def lean_capture_policy(
num_q_heads: int,
kv_group_num: int,
batch: int,
is_mla: bool = False,
) -> bool:
"""CUDA-graph capture-time bake decision for Lean decode.
During decode-graph capture ``seq_lens`` are set to the fill value (1), so the
seq-len based :func:`lean_decode_seqlen_gate` always sees ``avg_len == 1`` and returns
``False`` -- baking the *standard* kernel into every captured graph. Since the default
(auto) path replays those captured graphs, keying the bake on ``seq_lens_sum`` makes
Lean a no-op under CUDA graphs. But Lean's *fixed* 512-CTA persistent grid derives its
work schedule on-device from ``kv_indptr`` read at replay (work-stealing), so a baked
Lean graph still adapts to the real per-step raggedness. The bake decision therefore
keys only on capture-time-known signals -- ``batch``, the standard kernel's head-tile
parallelism (``tiles``), and ``is_mla`` -- calibrated to the realistic (ragged) regime.
Thresholds from ``CALIBRATION.md`` (MI355X, triton 3.7.0, CUDA graphs on, seed 42):
* MLA (``is_mla``): bake at ``batch >= 8``. ``b1`` is a catastrophic 0.40-0.55x loss
(the ~128-query-head decode already saturates the CUs, so Lean's fixed grid is pure
overhead), while ``b >= 8`` is uniform-parity and a 1.09-1.18x ragged win (p99 ITL
1.24-1.99x lower). This replaces the former blanket ``is_mla -> off``, which was based
on unrepresentative uniform ``b1`` data.
* GQA/MHA (``tiles >= 4``): bake at ``batch >= 16`` -- the unconditional-win boundary at
both 16K and 64K (up to 1.49x on ragged, 1.05-1.13x even on uniform). ``batch < 16``
is context-split (``b1`` wins at 64K but ``b1-4`` lose at 16K) and cannot be decided
from batch alone, so it is left to the eager :func:`lean_decode_seqlen_gate`.
* Heavy TP shard (``tiles < 4``, e.g. Llama-70B @TP=8): never bake. Not calibrated for
capture and known to regress ~4x at 32K; its rare long-context win still activates via
the eager seq-len gate (131072 base).
"""
if batch <= 0:
return False
if is_mla:
return batch >= 8
if _lean_head_tiles(num_q_heads, kv_group_num) < 4:
return False
return batch >= 16
def lean_decode_seqlen_gate(
num_q_heads: int,
kv_group_num: int,
batch: int,
seq_lens_sum: Optional[int],
is_mla: bool = False,
) -> bool:
"""Cheap host-side pre-gate for Lean decode (no GPU sync).
Lean Attention only beats the standard decode kernel for long-enough sequences; for
short context it both loses and would pay a ``kv_indptr[batch].item()`` host-sync in
:func:`decode_attention_fwd` just to discover it should fall back. The attention backend
calls this first, using host-side metadata it already has (``num_q_heads``,
``kv_group_num``, ``seq_lens_sum``, ``batch``), so short-context decode skips Lean
entirely without a sync.
What actually drives the Lean-vs-SplitK crossover is how well the standard grouped
kernel already fills the device, i.e. its query-head **parallelism**, not ``kv_group_num``.
The standard kernel launches ``tiles = ceil(num_q_heads / min(16, kv_group_num))``
head-tile programs per (sequence, kv-split); when ``tiles`` is large it saturates the CUs
at short context and Lean wins only much later, while with few query heads per GPU (heavy
tensor-parallel shards) it under-fills and Lean needs a long context to amortise its
fixed persistent-grid overhead. Keying the threshold on ``kv_group_num`` alone mispredicts
this badly: e.g. Llama-3-70B at TP=8 (8 query heads/GPU, ``kv_group_num`` still 8) is 4x
SLOWER under Lean at 32K, yet the old gate enabled it there. So we tier the base threshold
on ``tiles`` instead. Thresholds are the crossovers measured by ``benchmark/lean_gate_sweep.py``
on MI355X (256 CUs); they should scale with the device CU count on other GPUs.
MLA layers (``is_mla``, i.e. ``qk_head_dim != v_head_dim``) are gated on ``batch`` rather
than average length: their Lean win is driven by batch raggedness (work-stealing across
mixed-length requests), not context. Calibration shows ``b1`` loses hard (~0.40-0.55x: the
~128-query-head decode already saturates the CUs at batch 1, so Lean's fixed persistent grid
is pure overhead) while ``b >= 8`` is uniform-parity and a ragged win, so MLA enables at
``batch >= 8`` above a small length floor. This replaces the former blanket ``is_mla -> off``
(which was based on unrepresentative uniform ``b1`` data). In practice MLA models serve under
CUDA graphs, where :func:`lean_capture_policy` -- not this seq-len gate -- makes the decision.
The thresholds are set for the END-TO-END crossover, which is LATER than the isolated
kernel crossover: Lean's decode kernel has a nearly flat per-call cost, so even after the
standard kernel's attention becomes slower the *whole decode step* only turns over once
the standard attention has grown enough to clear Lean's flat floor. Measured end-to-end on
several GQA models, the crossover clusters at ~56-64K largely independent of the exact tile
count: Qwen2.5-7B (tiles=4) 0.84x@32K, 1.11x@64K, 1.88x@128K; Llama-3.1-8B (tiles=8)
1.06x@64K, 1.83x@128K; Ministral-8B (tiles=8) 0.86x@32K. So grouped decode uses a single
64K base and only heavy tensor-parallel shards with very few query-head tiles (tiles<4,
e.g. Llama-70B @TP=8, whose kernel crossover is already ~128K) push it to 128K. MHA (many
tiles) uses a lower 16K base (kernel crossover ~8K).
Thresholds relax as the batch grows, since more concurrent requests fill the persistent
grid at shorter lengths, and Lean is never enabled below a floor of 4K average tokens,
keeping the workload clear of the degenerate tiny-tile regime. The relaxation rate is
tier-dependent, from a saturated batch sweep on MI355X (range-ratio 0.25 ragged, batch =
concurrency, num_prompts>=6*batch):
* ``tiles >= 4`` (GQA/MHA): divisor ``batch // 2``. Measured E2E (throughput / median ITL)
confirms Lean wins well below the old ``batch // 4`` threshold once the batch fills the
grid. Qwen2.5-7B (tiles=4) @ batch: b4 0.997x/0.97x (neutral -> keep off), b8 1.05x/1.26x,
b12 1.13x/1.40x, b16 1.20x/1.84x, b32 1.28x/2.31x @ ~18.75K; and @ batch 16 it already
wins by ~7.5K avg (1.13x/1.39x). Llama-3.1-8B (tiles=8) @ batch 16 wins at every context
7.5K->30K (1.27-1.31x thrpt, 1.6-2.2x ITL). ``batch // 2`` enables from batch 8 @ ~18K
and batch 16 @ ~8K while keeping batch 4 conservative (32K threshold, correctly off at
18.75K where Lean is neutral).
* ``tiles < 4`` (heavy TP shard): keeps the conservative ``batch // 4``. Its E2E win needs
very long context (Llama-70B @TP=8 was 4x SLOWER at 32K); the isolated kernel can win at
high batch/long context but that does not survive the MoE + TP-all-reduce full step, and
a single-GPU microbench cannot replicate it, so this tier stays protected.
"""
if batch <= 0:
return False
# No CPU length mirror (e.g. gpu-only batches, or the EAGLE draft runner, which leaves
# seq_lens_sum unset): we cannot judge context length, so fall back to the standard kernel.
if seq_lens_sum is None:
return False
avg_len = seq_lens_sum / batch
if is_mla:
# Gate on batch (raggedness proxy), not average length; b1 is a hard loss, b>=8
# is parity-uniform / ragged-win. The 4K floor keeps degenerate tiny workloads off.
return batch >= 8 and avg_len >= 4096
tiles = _lean_head_tiles(num_q_heads, kv_group_num)
if tiles >= 16:
base = (
16384 # MHA / many query heads: standard kernel fills late, Lean wins early
)
elif tiles >= 4:
base = 65536 # typical GQA: measured E2E crossover ~56-64K
else:
base = (
131072 # few query heads/GPU (heavy TP shard): Lean needs very long context
)
# Heavy TP shards (tiles<4) relax slowly (batch//4); GQA/MHA relax at batch//2, matching
# the measured saturated-batch crossovers (see docstring).
div = batch // 2 if tiles >= 4 else batch // 4
threshold = max(4096, base // max(1, div))
return avg_len >= threshold
def _should_use_lean_decode(
enable_lean: Optional[bool],
logit_cap: float,
sinks,
xai_temperature_len: int,
score_mod,
) -> bool:
"""Decide whether the Work-Centric (Lean) Attention decode kernel may be used.
``enable_lean`` is the resolved activation flag passed by the caller:
* ``False`` — never use Lean Attention.
* ``True`` — use Lean Attention (the caller has already decided it is appropriate).
* ``None`` — do NOT self-enable here. Lean is only beneficial for long sequences and
its persistent-grid schedule misbehaves on tiny workloads, but this function has no
cheap way to know the sequence length (reading it would force a host sync that breaks
CUDA-graph capture). The attention backend resolves ``None`` to ``True``/``False`` via
:func:`lean_decode_seqlen_gate` using host-side metadata before calling in, so a
``None`` that reaches here (e.g. a direct call) conservatively means "off".
Regardless of the override, Lean Attention is only eligible when the request uses
none of the features the kernel does not implement. The kernel supports MHA, GQA, and
MLA (rope split), but ignores logit capping, attention sinks, xAI temperature scaling,
and score modification, so we fall back to the standard kernel whenever any of those
are requested rather than silently returning wrong results.
"""
if not enable_lean: # False or None
return False
if logit_cap and logit_cap > 0:
return False
if sinks is not None:
return False
if xai_temperature_len and xai_temperature_len > 0:
return False
if score_mod is not None:
return False
return True
def _lean_decode_launch_params(num_kv_heads, kv_group_num):
"""Lean decode launch parameters that depend only on shape (no seqlen, no sync).
Returns ``(total_programs, XCD_REMAP, NUM_XCDS)``. ``total_programs`` is the fixed
persistent-grid size (2× device CU count for better occupancy, rounded to a whole
number of XCDs when the XCD remap is active). The per-call tile schedule is computed
inside the kernel from ``kv_indptr``. Shared by :func:`decode_attention_fwd` and the
test so the grid/XCD decision stays in sync with the kernel.
"""
num_head_blocks = (kv_group_num + _LEAN_BLOCK_M - 1) // _LEAN_BLOCK_M
# XCD remap for ROCm only when rows are one-per-kv-head and divisible by 8.
XCD_REMAP = (num_kv_heads % 8 == 0 and num_head_blocks == 1) if _is_hip else False
NUM_XCDS = 8 if XCD_REMAP else 1
# Grid = round(CU_count * multiplier); multiplier defaults to 1.0 (one CTA per CU) and
# is overridable via SGLANG_FORCE_LEAN_GRID_CU_MULT for grid A/B tuning without a rebuild.
total_programs = max(
1, round(_lean_num_cus() * envs.SGLANG_FORCE_LEAN_GRID_CU_MULT.get())
)
if XCD_REMAP:
# The XCD remap requires the grid to be a whole number of XCDs.
total_programs = max((total_programs // NUM_XCDS) * NUM_XCDS, NUM_XCDS)
return total_programs, XCD_REMAP, NUM_XCDS
def _decode_lean_attention_fwd(
q,
k_buffer,
v_buffer,
o,
kv_indptr,
kv_indices,
total_programs,
sm_scale, # already folded with k_scale by the caller (matches the standard kernel)
v_scale,
XCD_REMAP,
NUM_XCDS,
Mp,
Lp,
Op,
locks,
page_size=1,
):
"""Wrapper for Lean Attention kernel.
``total_programs`` is the fixed persistent-grid size (2× device CU count). The kernel
derives its own tile schedule from ``kv_indptr`` on-device, so no host sync is needed and
the launch is CUDA-graph capturable. ``Mp``, ``Lp``, ``Op``, ``locks`` are pre-allocated
persistent-grid partial-result buffers reused across decode steps. ``page_size`` selects
the KV address math: 1 for a contiguous ``[N, head, dim]`` buffer, >1 for a paged
``[num_pages, page_size, head, dim]`` buffer (strides via ``_extract_kv_strides``).
"""
batch, head_num = q.shape[0], q.shape[1]
# head_num lives at dim -2 for both the 3-D [N, head, dim] and 4-D paged
# [num_pages, page_size, head, dim] layouts.
num_kv_heads = k_buffer.shape[-2]
Lk = k_buffer.shape[-1]
Lv = v_buffer.shape[-1]
kv_group_num = head_num // num_kv_heads
# MLA rope split: K carries an extra positional-encoding block (Lk > Lv).
if Lk == 576:
BLOCK_DMODEL, BLOCK_DPE = 512, 64
elif Lk == 288:
BLOCK_DMODEL, BLOCK_DPE = 256, 32
else:
BLOCK_DMODEL, BLOCK_DPE = triton.next_power_of_2(Lk), 0
BLOCK_DV = triton.next_power_of_2(Lv)
BLOCK_M = _LEAN_BLOCK_M
BLOCK_N = _lean_decode_block_n(Lk)
# A kv group wider than BLOCK_M is processed as several head blocks; each
# (kv_head, head_block) pair is one scheduling "row".
num_head_blocks = (kv_group_num + BLOCK_M - 1) // BLOCK_M
num_rows = num_kv_heads * num_head_blocks
rows_per_xcd = num_rows // NUM_XCDS if XCD_REMAP else num_rows
xcd_programs = total_programs // NUM_XCDS if XCD_REMAP else total_programs
# Pre-allocated persistent-grid partial-result buffers (Mp, Lp, Op, locks) are passed
# in and reused across decode steps; they hold running softmax state for BLOCK_M query
# heads (one head block of a kv group) per program. Reset locks to zero each call.
locks.zero_()
# Prepare batch_num_block_n (cumulative tiles per sequence) over the active batch.
# seq_len[i] = kv_indptr[i+1] - kv_indptr[i]
seq_lens = (kv_indptr[1 : batch + 1] - kv_indptr[:batch]).to(
torch.int64
) # use int64 for safe arithmetic
tiles_per_batch = (seq_lens + (BLOCK_N - 1)) // BLOCK_N
batch_num_block_n = (
torch.cumsum(tiles_per_batch, dim=0).to(torch.int32).contiguous()
)
max_output_tile_cnt = math.ceil((head_num * batch) / total_programs) + 4
# Page-aware KV strides. For a 3-D buffer these synthesize page/tok strides so the
# PAGE_SIZE>1 math collapses to the contiguous slot address; for a 4-D paged buffer they
# come from the real page/token strides. (See _extract_kv_strides.)
k_bs, k_h, k_page, k_tok = _extract_kv_strides(k_buffer, page_size)
v_bs, v_h, v_page, v_tok = _extract_kv_strides(v_buffer, page_size)
_lean_attention_decode_kernel[(total_programs,)](
q,
k_buffer,
v_buffer,
Mp,
Lp,
Op,
o,
batch_num_block_n,
locks,
kv_indptr,
kv_indices,
sm_scale,
v_scale,
q.stride(0),
q.stride(1),
k_bs,
k_h,
k_page,
k_tok,
v_bs,
v_h,
v_page,
v_tok,
o.stride(0),
o.stride(1),
kv_group_num=kv_group_num,
NUM_HEAD_BLOCKS=num_head_blocks,
ROWS_PER_XCD=rows_per_xcd,
BLOCK_DMODEL=BLOCK_DMODEL,
BLOCK_DPE=BLOCK_DPE,
BLOCK_DV=BLOCK_DV,
BLOCK_M=BLOCK_M,
BLOCK_N=BLOCK_N,
PAGE_SIZE=page_size,
XCD_REMAP=XCD_REMAP,
NUM_XCDS=NUM_XCDS,
batch_size=batch,
total_programs=total_programs,
num_query_heads=head_num,
num_rows=num_rows,
xcd_programs=xcd_programs,
max_output_tile_cnt=max_output_tile_cnt,
Lk=Lk,
Lv=Lv,
num_warps=4,
num_stages=2,
)
python/sglang/srt/environ.py
+13
View File
@@ -960,6 +960,19 @@ class Envs:
# ===================================================================
SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False)
SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE = EnvBool(False)
# A-B kill-switch for Work-Centric (Lean) Attention. When True, forces the
# standard Triton decode kernel even if --enable-lean-attention or the auto-gate
# would select Lean. Used to isolate the Lean kernel in benchmarks.
SGLANG_DISABLE_LEAN_ATTENTION = EnvBool(False)
# Persistent-grid size multiplier for the Lean decode kernel:
# total_programs = round(device_CU_count * this). Default 1.0 (one CTA per CU), which
# maximizes KV work-tiles per CTA and minimizes the cross-CTA combine/atomic reduction.
# Kernel + E2E A/B sweeps found 1.0 beats 2.0 across uniform and ragged configs on both
# MI300X (gfx942) and MI355X (gfx950) — 2.0 oversubscribed the CUs and regressed high-batch
# decode. Exposed as a knob (e.g. set 2.0) for grid A/B tuning without a rebuild.
SGLANG_FORCE_LEAN_GRID_CU_MULT = EnvFloat(1.0)
# Torch Compile
# Compact extend-attention query-tile grid: AMD/HIP-only optimization
# (parity with flash-attn's ragged-aware launch). The feature checks _is_hip
# explicitly in code; this env var allows override (0=force off, 1=force on).
python/sglang/srt/layers/attention/triton_backend.py
+119
View File
@@ -136,6 +136,11 @@ class ForwardMetadata:
# PHYSICAL full-attn write target for the unified pool (eager: translated tensor;
# cuda-graph: capture-stable buffer view). None for non-unified pools.
out_cache_loc_full_physical: Optional[torch.Tensor] = None
# Lean decode (persistent-grid partial-result buffers)
lean_Mp: Optional[torch.Tensor] = None
lean_Lp: Optional[torch.Tensor] = None
lean_Op: Optional[torch.Tensor] = None
lean_locks: Optional[torch.Tensor] = None
class TritonAttnBackend(AttentionBackend):
@@ -155,7 +160,11 @@ class TritonAttnBackend(AttentionBackend):
):
# Lazy import to avoid the initialization of cuda context
from sglang.kernels.ops.attention.decode_attention import (
_LEAN_BLOCK_M,
_lean_decode_launch_params,
decode_attention_fwd,
lean_capture_policy,
lean_decode_seqlen_gate,
)
from sglang.kernels.ops.attention.extend_attention import (
build_unified_kv_indices,
@@ -172,6 +181,11 @@ class TritonAttnBackend(AttentionBackend):
super().__init__()
self.decode_attention_fwd = torch.compiler.disable(decode_attention_fwd)
# Work-Centric (Lean) Attention activation. None => auto-gate from host-side
# seqlen metadata in forward_decode; True/False => explicit override.
self.enable_lean_attention = model_runner.server_args.enable_lean_attention
self._lean_decode_seqlen_gate = lean_decode_seqlen_gate
self._lean_capture_policy = lean_capture_policy
self.extend_attention_fwd = torch.compiler.disable(extend_attention_fwd)
self.extend_attention_fwd_unified = torch.compiler.disable(
extend_attention_fwd_unified
@@ -256,6 +270,14 @@ class TritonAttnBackend(AttentionBackend):
self.max_context_len = model_runner.model_config.context_len
self.device = model_runner.device
self.device_core_count = get_device_core_count(model_runner.gpu_id)
# Lean decode persistent-grid size (depends only on head architecture).
kv_group_num = self.num_head // self.num_kv_head
self.lean_total_programs, _, _ = _lean_decode_launch_params(
self.num_kv_head, kv_group_num
)
# BLOCK_M for Lean partial-result buffers; kept as an attribute so the
# cuda-graph / eager buffer allocators (separate methods) can size them.
self.lean_block_m = _LEAN_BLOCK_M
self.static_kv_splits = get_bool_env_var(
"SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS", "false"
)
@@ -778,6 +800,9 @@ class TritonAttnBackend(AttentionBackend):
window_kv_offsets = None
swa_attn_logits = None
spec_info = forward_batch.spec_info
# Lean decode buffers are only allocated on the decode path below; default
# to None so the shared ForwardMetadata constructor works for extend/verify.
lean_Mp = lean_Lp = lean_Op = lean_locks = None
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None or spec_info.kv_indptr is None:
@@ -858,6 +883,26 @@ class TritonAttnBackend(AttentionBackend):
),
)
# Lean decode persistent-grid partial-result buffers.
lean_Mp = torch.empty(
(self.lean_total_programs, self.lean_block_m),
dtype=torch.float32,
device=self.device,
)
lean_Lp = torch.empty(
(self.lean_total_programs, self.lean_block_m),
dtype=torch.float32,
device=self.device,
)
lean_Op = torch.empty(
(self.lean_total_programs, self.lean_block_m, self.v_head_dim),
dtype=torch.float32,
device=self.device,
)
lean_locks = torch.zeros(
(self.lean_total_programs,), dtype=torch.int32, device=self.device
)
qo_indptr = None
custom_mask = None
mask_indptr = None
@@ -1014,6 +1059,10 @@ class TritonAttnBackend(AttentionBackend):
swa_attn_logits=swa_attn_logits,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
lean_Mp=lean_Mp,
lean_Lp=lean_Lp,
lean_Op=lean_Op,
lean_locks=lean_locks,
)
def init_cuda_graph_state(
@@ -1047,6 +1096,26 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
# Lean decode persistent-grid partial-result buffers (shared across all layers).
self.cuda_graph_lean_Mp = torch.zeros(
(self.lean_total_programs, self.lean_block_m),
dtype=torch.float32,
device=self.device,
)
self.cuda_graph_lean_Lp = torch.zeros(
(self.lean_total_programs, self.lean_block_m),
dtype=torch.float32,
device=self.device,
)
self.cuda_graph_lean_Op = torch.zeros(
(self.lean_total_programs, self.lean_block_m, self.v_head_dim),
dtype=torch.float32,
device=self.device,
)
self.cuda_graph_lean_locks = torch.zeros(
(self.lean_total_programs,), dtype=torch.int32, device=self.device
)
if cuda_graph_num_kv_splits_buf is None:
self.cuda_graph_num_kv_splits = torch.full(
(max_num_tokens,),
@@ -1157,6 +1226,10 @@ class TritonAttnBackend(AttentionBackend):
swa_attn_logits=self.cuda_graph_swa_attn_logits,
swa_out_cache_loc=swa_out_cache_loc,
out_cache_loc_full_physical=out_cache_loc_full_physical,
lean_Mp=self.cuda_graph_lean_Mp,
lean_Lp=self.cuda_graph_lean_Lp,
lean_Op=self.cuda_graph_lean_Op,
lean_locks=self.cuda_graph_lean_locks,
)
elif forward_mode.is_target_verify():
custom_mask = (
@@ -1879,6 +1952,42 @@ class TritonAttnBackend(AttentionBackend):
):
attn_logits = self.forward_metadata.swa_attn_logits
# Resolve Work-Centric (Lean) Attention activation. In auto mode (None) the decision
# depends on whether this forward is a CUDA-graph capture: during capture seq_lens are
# the fill value (1), so the seq-len gate would always bake the standard kernel and Lean
# would never activate on the default path. There we key the bake on capture-time-known
# signals (batch, head-tiles, is_mla) via lean_capture_policy -- Lean's fixed persistent
# grid still adapts to raggedness on-device at replay. In eager decode, real seq_lens
# are known, so lean_decode_seqlen_gate uses them. An explicit True/False override is
# respected; the SGLANG_DISABLE_LEAN_ATTENTION kill-switch forces the standard kernel.
from sglang.srt.environ import envs
from sglang.srt.model_executor.runner_utils.capture_mode import (
get_is_capture_mode,
)
if envs.SGLANG_DISABLE_LEAN_ATTENTION.get():
enable_lean = False
else:
enable_lean = self.enable_lean_attention
if enable_lean is None:
kv_group_num = layer.tp_q_head_num // layer.tp_k_head_num
is_mla = layer.qk_head_dim != layer.v_head_dim
if get_is_capture_mode():
enable_lean = self._lean_capture_policy(
layer.tp_q_head_num,
kv_group_num,
forward_batch.batch_size,
is_mla,
)
else:
enable_lean = self._lean_decode_seqlen_gate(
layer.tp_q_head_num,
kv_group_num,
forward_batch.batch_size,
forward_batch.seq_lens_sum,
is_mla,
)
if self.dcp_size > 1:
if score_mod is not None:
raise NotImplementedError(
@@ -1913,6 +2022,11 @@ class TritonAttnBackend(AttentionBackend):
logit_cap=logits_soft_cap,
sinks=sinks,
xai_temperature_len=layer.xai_temperature_len,
enable_lean=enable_lean,
lean_Mp=self.forward_metadata.lean_Mp,
lean_Lp=self.forward_metadata.lean_Lp,
lean_Op=self.forward_metadata.lean_Op,
lean_locks=self.forward_metadata.lean_locks,
)
local_lse = torch.logsumexp(
self.forward_metadata.attn_lse[
@@ -1945,6 +2059,11 @@ class TritonAttnBackend(AttentionBackend):
page_size=self.page_size,
score_mod=score_mod,
aux_tensors=aux_tensors,
enable_lean=enable_lean,
lean_Mp=self.forward_metadata.lean_Mp,
lean_Lp=self.forward_metadata.lean_Lp,
lean_Op=self.forward_metadata.lean_Op,
lean_locks=self.forward_metadata.lean_locks,
)
return o
python/sglang/srt/server_args.py
+5
View File
@@ -1666,6 +1666,11 @@ class ServerArgs:
),
NS("exec.kernel"),
] = None
enable_lean_attention: A[
Optional[bool],
"Enable Lean (Work-Centric) Attention decode kernel for long-context serving. When None (default), uses auto-gate that activates Lean for long contexts and falls back to standard kernel for short contexts. Set to True to force enable, False to force disable.",
NS("exec.kernel"),
] = None
prefill_attention_backend: A[
Optional[str],
Arg(
test/registered/kernels/test_lean_attention.py
+453
View File
@@ -0,0 +1,453 @@
"""CI gating test for Work-Centric (Lean) Attention decode kernel.
Lean is an opt-in, gated decode-attention kernel in the ROCm/AMD Triton backend
(``python/sglang/kernels/ops/attention/decode_attention.py``). Its core contract
is that it is **numerically identical** to the standard SplitK grouped kernel —
the auto-gate only decides *when* to use it for speed, never *whether* the output
is correct. This test locks in that contract so a future change to the kernel or
its launch/reduction path cannot silently regress correctness.
Two things are checked:
1. **Parity** — Lean output matches the standard SplitK kernel (cosine sim ~1.0)
across representative GQA head shapes / batches / contexts.
2. **Gate logic** — the eager ``lean_decode_seqlen_gate`` enables Lean in the
long-context / low-batch regime and keeps it off for short context, and the
CUDA-graph ``lean_capture_policy`` bakes Lean from capture-time signals (batch,
head-tiles, is_mla) since captured seq_lens are the fill value. MLA is gated on
batch (off at b1, on at b>=8), not blanket-off.
Correctness is triton-version-independent (only performance varies with the triton
build), so this makes a robust per-commit gate on MI35x hardware.
"""
import unittest
import torch
from sglang.kernels.ops.attention.decode_attention import (
_LEAN_BLOCK_M,
_lean_decode_launch_params,
decode_attention_fwd,
decode_attention_fwd_grouped,
lean_capture_policy,
lean_decode_seqlen_gate,
)
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
# Lean lives in the ROCm/AMD Triton backend and is tuned for gfx950 (MI35x),
# so gate it on the per-commit MI35x single-GPU suite. Correctness (not perf) is
# what this test asserts, which holds regardless of the triton build.
register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x")
# (name, H_Q, H_KV, head_dim) — the two GQA shapes validated in the PR benchmarks.
GQA_SHAPES = [
("qwen2.5-7b", 28, 4, 128),
("llama3.1-8b", 32, 8, 128),
]
MAX_KV_SPLITS = 8
def _run_pair(H_Q, H_KV, D, B, S, dev="cuda", dt=torch.float16, seed=0):
"""Run standard SplitK and Lean on the same inputs; return (o_std, o_lean)."""
torch.manual_seed(seed)
D_V = D
kv_group_num = H_Q // H_KV
sm = 1.0 / (D**0.5)
tot = B * S
total_programs, _, _ = _lean_decode_launch_params(H_KV, kv_group_num)
lean_Mp = torch.empty(
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
)
lean_Lp = torch.empty(
(total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev
)
lean_Op = torch.empty(
(total_programs, _LEAN_BLOCK_M, D_V), dtype=torch.float32, device=dev
)
lean_locks = torch.zeros((total_programs,), dtype=torch.int32, device=dev)
kv_indptr = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
kv_indices = torch.arange(0, tot, device=dev, dtype=torch.int32)
q = torch.randn(B, H_Q, D, dtype=dt, device=dev)
k = torch.randn(tot, H_KV, D, dtype=dt, device=dev)
v = torch.randn(tot, H_KV, D_V, dtype=dt, device=dev)
num_kv_splits = torch.full((B,), MAX_KV_SPLITS, dtype=torch.int32, device=dev)
attn_logits = torch.empty(
(B, H_Q, MAX_KV_SPLITS, D_V), dtype=torch.float32, device=dev
)
attn_lse = torch.empty((B, H_Q, MAX_KV_SPLITS), dtype=torch.float32, device=dev)
o_std = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
decode_attention_fwd_grouped(
q,
k,
v,
o_std,
kv_indptr,
kv_indices,
attn_logits,
attn_lse,
num_kv_splits,
MAX_KV_SPLITS,
sm,
1.0,
)
attn_logits2 = torch.empty_like(attn_logits)
attn_lse2 = torch.empty_like(attn_lse)
o_lean = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
decode_attention_fwd(
q,
k,
v,
o_lean,
kv_indptr,
kv_indices,
attn_logits2,
attn_lse2,
num_kv_splits,
MAX_KV_SPLITS,
sm,
1.0,
1.0,
enable_lean=True,
lean_Mp=lean_Mp,
lean_Lp=lean_Lp,
lean_Op=lean_Op,
lean_locks=lean_locks,
)
return o_std, o_lean
def _lean_scratch(H_KV, kv_group_num, D_V, dev):
total_programs, _, _ = _lean_decode_launch_params(H_KV, kv_group_num)
return (
torch.empty((total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev),
torch.empty((total_programs, _LEAN_BLOCK_M), dtype=torch.float32, device=dev),
torch.empty(
(total_programs, _LEAN_BLOCK_M, D_V), dtype=torch.float32, device=dev
),
torch.zeros((total_programs,), dtype=torch.int32, device=dev),
)
def _run_pair_fp8(H_Q, H_KV, D, B, S, fp8_dtype, dev="cuda", seed=0):
"""Standard vs Lean on **fp8** K/V with non-unit k_scale/v_scale.
Both arms go through the public ``decode_attention_fwd`` dispatch (enable_lean False/True),
which folds k_scale into sm_scale and applies v_scale — the exact production path. They share
the same fp8 inputs and dequant scales, so their outputs must agree (the fp8 quantization
error is identical for both); this guards that Lean's fp8 dtype handling matches the standard
kernel. Returns (o_std, o_lean).
"""
torch.manual_seed(seed)
D_V = D
kv_group_num = H_Q // H_KV
sm = 1.0 / (D**0.5)
tot = B * S
fp8_max = torch.finfo(fp8_dtype).max
q = torch.randn(B, H_Q, D, dtype=torch.float16, device=dev)
k_ref = torch.randn(tot, H_KV, D, dtype=torch.float32, device=dev)
v_ref = torch.randn(tot, H_KV, D_V, dtype=torch.float32, device=dev)
# Per-tensor symmetric quantization to fp8, mirroring how fp8 KV is stored + dequantized.
k_scale = (k_ref.abs().max() / fp8_max).item()
v_scale = (v_ref.abs().max() / fp8_max).item()
k = (k_ref / k_scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
v = (v_ref / v_scale).clamp(-fp8_max, fp8_max).to(fp8_dtype)
kv_indptr = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
kv_indices = torch.arange(0, tot, device=dev, dtype=torch.int32)
num_kv_splits = torch.full((B,), MAX_KV_SPLITS, dtype=torch.int32, device=dev)
def _call(enable_lean):
attn_logits = torch.empty(
(B, H_Q, MAX_KV_SPLITS, D_V), dtype=torch.float32, device=dev
)
attn_lse = torch.empty((B, H_Q, MAX_KV_SPLITS), dtype=torch.float32, device=dev)
o = torch.zeros(B, H_Q, D_V, dtype=torch.float16, device=dev)
mp, lp, op, locks = _lean_scratch(H_KV, kv_group_num, D_V, dev)
decode_attention_fwd(
q,
k,
v,
o,
kv_indptr,
kv_indices,
attn_logits,
attn_lse,
num_kv_splits,
MAX_KV_SPLITS,
sm,
k_scale,
v_scale,
enable_lean=enable_lean,
lean_Mp=mp,
lean_Lp=lp,
lean_Op=op,
lean_locks=locks,
)
return o
return _call(False), _call(True)
def _run_pair_paged(
H_Q, H_KV, D, B, S, page_size, dev="cuda", dt=torch.float16, seed=0
):
"""Standard vs Lean on a **paged** 4-D KV buffer ``[num_pages, page_size, head, dim]``.
The KV cache is stored in pages and addressed through scattered slot ids in ``kv_indices``
(a permutation), so the kernel's page-aware address math (``kv_loc // page_size`` /
``kv_loc % page_size``) is genuinely exercised — not the contiguous fast path. Both arms read
the identical buffer + indices, so their outputs must agree. Returns (o_std, o_lean).
"""
torch.manual_seed(seed)
D_V = D
kv_group_num = H_Q // H_KV
sm = 1.0 / (D**0.5)
tot = B * S
assert (
tot % page_size == 0
), "test setup: total tokens must be a multiple of page_size"
num_pages = tot // page_size
# 4-D paged KV buffers [num_pages, page_size, head, dim] (the shared-pool layout).
k = torch.randn(num_pages, page_size, H_KV, D, dtype=dt, device=dev)
v = torch.randn(num_pages, page_size, H_KV, D_V, dtype=dt, device=dev)
kv_indptr = torch.arange(0, (B + 1) * S, step=S, device=dev, dtype=torch.int32)
# Scatter slots across pages so page_id/tok_in_p vary within every BLOCK_N tile.
kv_indices = torch.randperm(tot, device=dev).to(torch.int32)
q = torch.randn(B, H_Q, D, dtype=dt, device=dev)
num_kv_splits = torch.full((B,), MAX_KV_SPLITS, dtype=torch.int32, device=dev)
def _call(enable_lean):
attn_logits = torch.empty(
(B, H_Q, MAX_KV_SPLITS, D_V), dtype=torch.float32, device=dev
)
attn_lse = torch.empty((B, H_Q, MAX_KV_SPLITS), dtype=torch.float32, device=dev)
o = torch.zeros(B, H_Q, D_V, dtype=dt, device=dev)
mp, lp, op, locks = _lean_scratch(H_KV, kv_group_num, D_V, dev)
decode_attention_fwd(
q,
k,
v,
o,
kv_indptr,
kv_indices,
attn_logits,
attn_lse,
num_kv_splits,
MAX_KV_SPLITS,
sm,
1.0,
1.0,
page_size=page_size,
enable_lean=enable_lean,
lean_Mp=mp,
lean_Lp=lp,
lean_Op=op,
lean_locks=locks,
)
return o
return _call(False), _call(True)
@unittest.skipUnless(torch.cuda.is_available(), "Lean decode kernel requires a GPU")
class TestLeanAttentionParity(CustomTestCase):
"""Lean must be numerically identical to the standard SplitK kernel."""
def test_parity_across_gqa_shapes(self):
# Contexts kept modest so CI stays fast; parity is context-independent.
for name, H_Q, H_KV, D in GQA_SHAPES:
for B in (1, 8):
for S in (8192, 32768):
with self.subTest(model=name, batch=B, ctx=S):
o_std, o_lean = _run_pair(H_Q, H_KV, D, B, S)
cos = torch.nn.functional.cosine_similarity(
o_lean.flatten().float(), o_std.flatten().float(), dim=0
).item()
self.assertGreater(
cos,
0.999,
f"{name} b={B} ctx={S}: Lean diverged from SplitK (cos={cos:.5f})",
)
# No NaN/Inf leaked from the persistent-grid reduction.
self.assertTrue(
torch.isfinite(o_lean).all(),
f"{name}: non-finite Lean output",
)
def test_fp8_kv_parity(self):
# Phase 2: Lean must handle fp8 KV cache the same way the standard kernel does
# (cast q->K.dtype for the MMA, fold k_scale into sm_scale, apply v_scale). Guards the
# regression where the Lean path crashed on fp8 K ("Unsupported rhs dtype fp8e4nv").
fp8_dtype = None
for name in ("float8_e4m3fn", "float8_e4m3fnuz"):
if hasattr(torch, name):
fp8_dtype = getattr(torch, name)
break
if fp8_dtype is None:
self.skipTest("no fp8 e4m3 dtype available in this torch build")
for name, H_Q, H_KV, D in GQA_SHAPES:
for B in (1, 8):
for S in (8192, 32768):
with self.subTest(model=name, batch=B, ctx=S, dtype=str(fp8_dtype)):
o_std, o_lean = _run_pair_fp8(H_Q, H_KV, D, B, S, fp8_dtype)
self.assertTrue(
torch.isfinite(o_lean).all(),
f"{name}: non-finite Lean fp8 output",
)
cos = torch.nn.functional.cosine_similarity(
o_lean.flatten().float(), o_std.flatten().float(), dim=0
).item()
self.assertGreater(
cos,
0.99,
f"{name} b={B} ctx={S}: Lean fp8 diverged from SplitK (cos={cos:.5f})",
)
def test_paged_kv_parity(self):
# Lean must read a paged 4-D KV buffer the same way the standard kernel does. Guards
# the page-aware address math (kv_loc // page_size, kv_loc % page_size); a regression
# to the contiguous-only form would scramble reads and drop cos well below 1.
for name, H_Q, H_KV, D in GQA_SHAPES:
for page_size in (16, 64):
with self.subTest(model=name, page_size=page_size):
o_std, o_lean = _run_pair_paged(
H_Q, H_KV, D, B=2, S=8192, page_size=page_size
)
self.assertTrue(
torch.isfinite(o_lean).all(),
f"{name} ps={page_size}: non-finite Lean paged output",
)
cos = torch.nn.functional.cosine_similarity(
o_lean.flatten().float(), o_std.flatten().float(), dim=0
).item()
self.assertGreater(
cos,
0.999,
f"{name} ps={page_size}: Lean paged diverged from SplitK (cos={cos:.5f})",
)
@unittest.skipUnless(
torch.cuda.is_available(), "gate is exercised alongside the kernel path"
)
class TestLeanSeqlenGate(CustomTestCase):
"""The auto-gate must enable Lean in its win region and stay off elsewhere."""
def test_gate_enables_long_context_low_batch(self):
# Qwen GQA (28Q/4KV): long context at batch 1 is squarely Lean's win region.
H_Q, kv_group = 28, 7
self.assertTrue(
lean_decode_seqlen_gate(
H_Q, kv_group, batch=1, seq_lens_sum=131072, is_mla=False
),
"gate should enable Lean for batch=1 @ 128K",
)
def test_gate_off_for_short_context(self):
H_Q, kv_group = 28, 7
self.assertFalse(
lean_decode_seqlen_gate(
H_Q, kv_group, batch=1, seq_lens_sum=2048, is_mla=False
),
"gate should keep Lean off for batch=1 @ 2K (standard kernel wins)",
)
def test_gate_mla_batch_threshold(self):
# MLA is gated on batch, not a blanket off: b1 loses hard (CU-saturated), b>=8
# is parity/ragged-win. The eager gate must reflect that boundary. (Guards against
# both a regression to the old blanket-off and to an always-on for MLA.)
H_Q, kv_group = 128, 128
self.assertFalse(
lean_decode_seqlen_gate(
H_Q, kv_group, batch=1, seq_lens_sum=131072, is_mla=True
),
"MLA at batch=1 is a catastrophic loss; gate must stay off",
)
self.assertTrue(
lean_decode_seqlen_gate(
H_Q, kv_group, batch=8, seq_lens_sum=8 * 65536, is_mla=True
),
"MLA at batch>=8 with long context is a win; gate must enable",
)
def test_gate_off_when_seq_lens_sum_missing(self):
# The EAGLE draft runner (and gpu-only batches) call decode without a CPU length
# mirror, so seq_lens_sum is None. The gate must fall back to the standard kernel
# instead of dividing None by batch (which raised TypeError and crashed the
# scheduler under EAGLE3 speculative decoding).
H_Q, kv_group = 28, 7
self.assertFalse(
lean_decode_seqlen_gate(
H_Q, kv_group, batch=8, seq_lens_sum=None, is_mla=False
),
"gate must return False (not raise) when seq_lens_sum is None",
)
def test_gate_threshold_falls_with_batch(self):
# The crossover context falls as batch grows: a context that is below the
# single-request threshold should still enable Lean at higher batch.
H_Q, kv_group = 28, 7
ctx = 32768
low_batch = lean_decode_seqlen_gate(
H_Q, kv_group, batch=1, seq_lens_sum=ctx, is_mla=False
)
high_batch = lean_decode_seqlen_gate(
H_Q, kv_group, batch=8, seq_lens_sum=ctx * 8, is_mla=False
)
# At batch 8 the same per-request context should be at least as likely to enable Lean.
self.assertTrue(
high_batch or not low_batch,
"gate batch relaxation is inconsistent (higher batch should not be stricter)",
)
class TestLeanCapturePolicy(CustomTestCase):
"""The CUDA-graph capture-time bake policy keys on (tiles, is_mla, batch) only —
captured seq_lens are the fill value, so it cannot use context. These pin the
calibrated thresholds (CALIBRATION.md): a threshold drift or a degraded predicate
(always-on / always-off) turns the corresponding case red."""
def test_gqa_bakes_at_and_above_batch_16(self):
# Qwen GQA (28Q/4KV -> tiles=4): unconditional-win boundary is batch>=16.
H_Q, kv_group = 28, 7
self.assertFalse(
lean_capture_policy(H_Q, kv_group, batch=8, is_mla=False),
"GQA capture must not bake at batch=8 (context-split / uniform-short regresses)",
)
self.assertTrue(
lean_capture_policy(H_Q, kv_group, batch=16, is_mla=False),
"GQA capture must bake at batch>=16 (unconditional win)",
)
def test_mla_bakes_at_and_above_batch_8_never_at_low_batch(self):
# MLA (128Q/128KV): b1 is a catastrophic loss, b>=8 is parity/ragged-win.
H_Q, kv_group = 128, 128
self.assertFalse(
lean_capture_policy(H_Q, kv_group, batch=1, is_mla=True),
"MLA capture must never bake at batch=1 (0.4-0.55x loss)",
)
self.assertTrue(
lean_capture_policy(H_Q, kv_group, batch=8, is_mla=True),
"MLA capture must bake at batch>=8",
)
def test_heavy_tp_shard_never_bakes(self):
# tiles<4 (e.g. Llama-70B @TP=8: 8 query heads/GPU, kv_group=1 -> tiles=8?) —
# use a genuine heavy shard: 2 query heads, kv_group=1 -> tiles=2 (<4). Known ~4x
# regression at 32K, not calibrated for capture -> never bake even at high batch.
self.assertFalse(
lean_capture_policy(2, 1, batch=32, is_mla=False),
"heavy TP shard (tiles<4) must never bake under capture",
)
if __name__ == "__main__":
unittest.main(verbosity=3)