diff --git a/benchmark/lean_kernel_sweep.py b/benchmark/lean_kernel_sweep.py new file mode 100755 index 000000000..f6a6679ac --- /dev/null +++ b/benchmark/lean_kernel_sweep.py @@ -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() diff --git a/python/sglang/kernels/ops/attention/decode_attention.py b/python/sglang/kernels/ops/attention/decode_attention.py index 696990895..3884e1ce4 100644 --- a/python/sglang/kernels/ops/attention/decode_attention.py +++ b/python/sglang/kernels/ops/attention/decode_attention.py @@ -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, + ) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c84c24e2b..716e80cff 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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). diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 39c45dd9d..018157b50 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 91005b75b..9539c7a8b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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( diff --git a/test/registered/kernels/test_lean_attention.py b/test/registered/kernels/test_lean_attention.py new file mode 100644 index 000000000..73e5124d4 --- /dev/null +++ b/test/registered/kernels/test_lean_attention.py @@ -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)