[valechen]

Motivation

The standard flash-decode ("SplitK") attention kernel parallelizes decode by splitting each sequence's KV cache into chunks and assigning chunks to GPU compute units (CUs). While effective at short context, it leaves performance on the table in two common serving regimes:

• Long-context, low-batch — std's grid is batch × head-tiles × num_splits with a fixed split count, so a single (or few) long sequence cannot generate enough KV-split work to fill all CUs: many CUs sit idle while the few launched programs each grind through a long serial KV loop, and decode-step latency grows roughly linearly with context, making the workload occupancy-starved, not bandwidth-bound. As context windows reach 64K–128K+, attention dominates the
  decode step here.
• Ragged batches — with mixed sequence lengths, SplitK cannot rebalance a long request's work onto the CUs freed by short requests, so a few stragglers set the step latency.

Work-Centric Attention (WCA / "Lean") launches a fixed persistent grid sized to the device CU count and derives its KV-tile schedule on-device from kv_indptr, cooperatively splitting long sequences across CTAs (work-stealing) so no CU idles. Decode-attention latency stays nearly flat as context grows (1.06–1.30× E2E throughput even on uniform batches) and absorbs raggedness for free (up to 1.52× throughput, 3.62× lower ITL). Lean is numerically identical to the standard kernel
and ships behind a conservative auto-gate that enables it only where it wins — including by default under CUDA graphs — so existing workloads are unaffected.

Validated on AMD MI355X (gfx950) across GQA (Qwen2.5-7B, Llama-3.1-8B) and MLA (DeepSeek-R1-MXFP4, TP=4, fp8 KV). Also re-validated on AMD MI300X (gfx942) at the same 1×CU grid across GQA (Qwen2.5-7B, Llama-3.1-8B) and MLA (proxied with DeepSeek-V2-Lite, fp8 KV, since fp4/MXFP4 GEMM is gfx950-only).

Modifications

A minimal, mostly-additive change (6 files, ~+1540 lines) on the ROCm/Triton backend. Almost all of it is new code; the only edits to existing paths are the forward_decode dispatch hook and the decode_attention_fwd signature. The standard decode path is untouched when Lean is inactive.

File	Change
python/sglang/kernels/ops/attention/decode_attention.py (+822)	The Lean kernel (_lean_attention_decode_kernel) and its launch/dispatch helpers. decode_attention_fwd gains an enable_lean parameter and early dispatch to Lean. Supports bf16/fp16 and fp8 KV (fold k_scale into sm_scale, cast q → K.dtype for the fp8×fp8 MMA, dequantize rope-split kpe, apply v_scale), and both contiguous and paged KV (kv_loc // page_size / % page_size addressing). Auto-gate: lean_decode_seqlen_gate (eager) and lean_capture_policy (CUDA-graph capture).
test/registered/kernels/test_lean_attention.py (+440)	CI gating test on the MI35x AMD per-commit suite (stage-b-test-1-gpu-small-amd-mi35x): parity vs standard SplitK for bf16 (cos > 0.999), fp8 (cos > 0.99), and paged KV (cos > 0.999), plus gate/policy logic.
python/sglang/srt/layers/attention/triton_backend.py (+119)	Resolve enable_lean_attention in forward_decode (explicit override → auto-gate → kill-switch); allocate the Lean scratch buffers (Mp/Lp/Op/locks) for the eager and CUDA-graph paths.
benchmark/lean_kernel_sweep.py (+150)	Kernel-level Lean-vs-standard (SplitK) sweep: both GQA head configs × full batch/context grid, per-call latency, speedup, cosine parity, and eager-gate decision; writes a CSV.
python/sglang/srt/server_args.py (+4)	--enable-lean-attention flag (Optional[bool], default None = auto-gate).
python/sglang/srt/environ.py (+11)	Two env vars: SGLANG_DISABLE_LEAN_ATTENTION kill-switch (force the standard kernel) and SGLANG_FORCE_LEAN_GRID_CU_MULT (persistent-grid size multiplier, default 1.0 = one CTA per CU; a knob for grid A/B without a rebuild).

Design:

• Persistent grid. Launches total_programs = device CU count CTAs — one CTA per CU (rounded to whole XCDs under the ROCm XCD remap). One CTA per CU maximizes the KV work-tiles handled per CTA, so each CTA's cross-CTA combine/atomic reduction amortizes more useful work; oversubscribing (e.g. 2 CTAs/CU) halves tiles-per-CTA and inflates that reduction, regressing high-batch decode (measured on MI300X/gfx942 and confirmed on MI355X/gfx950). The grid is independent of batch×heads×splits, so it is CUDA-graph capturable and a single captured grid replays across batch sizes.
• On-device work-stealing. Each CTA computes its slice of the global KV-tile work list from kv_indptr; long sequences are cooperatively split across CTAs. The schedule is derived from kv_indptr at replay, so a captured graph adapts to the real per-step sequence lengths and raggedness — no host sync, no per-step relaunch.
• Auto-gate. The default (None) resolves to a concrete decision before launch; explicit True/False overrides it. Because Lean is a win only in specific regimes, the gate operates in two modes:
  • Eager decode, where real seq_lens are known: lean_decode_seqlen_gate keys a length threshold on query-head parallelism (tiles = ceil(num_q_heads / min(16, kv_group)), i.e. how well the standard kernel already fills the CUs), relaxed as batch grows.
  • CUDA-graph capture, where per-step lengths are not yet known: lean_capture_policy makes the bake decision from capture-time-known signals (batch, tiles, is_mla) — bake GQA/MHA at batch ≥ 16, MLA at batch ≥ 8, and never for heavy TP shards (tiles < 4). Lean's fixed grid then adapts to the actual lengths at replay.
• Shape/dtype coverage. MHA, GQA, and MLA (rope-split qk_head_dim != v_head_dim); bf16/fp16/fp8 KV; contiguous and paged KV. Falls back to the standard kernel for unimplemented features (logit capping, attention sinks, xAI temperature, score modification), so results are never silently wrong.

Usage:

• Default (None): auto-gate — Lean turns on automatically where it wins (including by default under CUDA graphs). No flag needed.
• Force on: --enable-lean-attention (with --attention-backend triton).
• Force off: --enable-lean-attention=False, or SGLANG_DISABLE_LEAN_ATTENTION=1.

Accuracy Tests

Test configuration:

Hardware	1× AMD Instinct MI355X (256 CUs, gfx950, 288 GB HBM3E)
Software	ROCm 7.2.0, PyTorch 2.9.1+rocm7.2.0, Triton 3.7.0 (AMD wheel)
Models (head shapes)	Qwen2.5-7B-Instruct (GQA 28Q/4KV, D=128) · Llama-3.1-8B-Instruct (GQA 32Q/8KV, D=128)
Method	Isolated decode-attention call; compare Lean vs standard SplitK output over three KV formats — bf16, fp8 (e4m3, per-tensor scales), and paged (4-D buffer, scattered slots)
Metric	Cosine similarity of the two output tensors

Command (the registered CI gating test):

python3 -m pytest test/registered/kernels/test_lean_attention.py -v
# Per shape: standard = decode_attention_fwd(..., enable_lean=False)
#            lean     = decode_attention_fwd(..., enable_lean=True, lean_Mp/Lp/Op/locks=...)
#            cos      = F.cosine_similarity(o_lean.flatten().float(), o_std.flatten().float(), dim=0)

Results. Lean matches the standard kernel across every tested configuration:

KV format	Configs	Result
bf16/fp16	59 cells (2 models × 6 batches × 5 contexts, less Llama b32×128K)	cosine = 1.0000 (bit-identical)
fp8 e4m3 (per-tensor scale)	2 models × {b1,b8} × {8K,32K}	cosine > 0.99 (fp8 quantization noise only)
paged (4-D, scattered slots)	2 models × page_size {16, 64}	cosine > 0.999

Every measured configuration is bit-identical in bf16 (the one unmeasured cell, Llama batch-32 × 128K, exceeds the microbenchmark's single contiguous KV tensor — it is a harness limit, not a kernel result). The committed CI gating test (test_lean_attention.py) reproduces the bf16, fp8, and paged parity checks on the MI35x per-commit suite. Because output is numerically equivalent, downstream generation quality is unchanged.

Speed Tests and Profiling

Unless noted otherwise, the tables below are on 1× MI355X (gfx950); the dedicated MI300X (gfx942) subsection re-runs the same suites on that arch. Both: ROCm 7.2.0, PyTorch2.9.1+rocm7.2.0, Triton 3.7.0 (AMD wheel; the production ROCm wheel — the stock rocm7.2 image's Triton 3.5.1 inverts the crossover and must not be used for Lean benchmarking). The baseline is the Triton backend's own standard decode kernel, not AITER: both arms use --attention-backend triton
and dispatch through the same decode_attention_fwd, differing only in the enable_lean branch, so this is a clean Lean-vs-standard A/B.

Kernel-level (isolated decode attention, speedup = standard ÷ lean)

Direct kernel call, fp16, random Q/K/V, page_size=1; batch ∈ {1,2,4,8,16,32} × context ∈ {8K,16K,32K,64K,128K}. Command: SGLANG_FORCE_LEAN_GRID_CU_MULT=1.0 python3 benchmark/lean_kernel_sweep.py

Qwen2.5-7B (28Q/4KV):

batch	8K	16K	32K	64K	128K
1	0.46×	0.92×	1.86×	3.30×	5.33×
8	0.59×	1.04×	1.28×	1.59×	1.95×
32	1.78×	1.99×	2.21×	1.74×	1.77×

Llama-3.1-8B (32Q/8KV):

batch	8K	16K	32K	64K	128K
1	0.79×	1.47×	2.41×	3.56×	5.92×
8	1.46×	1.82×	1.97×	2.18×	1.71×
32	2.13×	2.29×	1.74×	1.73×	—

Bold = Lean is faster at the kernel level (speedup ≥ 1×). Peak 5.33× (Qwen) / 5.92× (Llama) at batch=1 / 128K; speedup grows superlinearly with context (SplitK's idle-CU penalty scales with length, Lean's stays flat) and the crossover falls with batch. The production auto-gate turns Lean on around the end-to-end break-even (later than the isolated-kernel crossover), so it enables the long-context / larger-batch cells and keeps the standard kernel for short-context low-batch. (The batch-32 × 128K Llama cell exceeds the microbenchmark's single contiguous KV tensor and is not
measured; real serving uses a paged KV pool and is unaffected — see the E2E tables.)

End-to-end serving — GQA

Full-model serving, closed-loop / fixed-batch (num-prompts == max-concurrency == reqs, --request-rate inf), CUDA graphs on (--cuda-graph-bs 1 2 4 8 16 32), --seed 42 (both arms see the identical batch). --random-range-ratio r samples per-request prompt and output lengths in [r·ctx, ctx], so ratio < 1 = ragged. Rows include batch 1 to show the low-batch long-context win and the auto-gate's enabled region (GQA batch ≥ 16); below b16 the standard kernel runs by default.

Columns: reqs = concurrent requests; KV mean / KV range = mean and min–max prompt (KV) length across the batch; dec/req = mean decode (output) tokens per request; dec iters = total decode tokens measured; tps× = Lean ÷ baseline output throughput; ITL× / p99× = baseline ÷ Lean median / p99 inter-token latency (>1 = Lean faster).

model	B	ratio	reqs	KV mean	KV range	dec/req	dec iters	tps×	ITL×	p99×
Qwen2.5-7B	1	1.0	1	65536	65K (uniform)	512	512	1.27	1.63	1.63
Qwen2.5-7B	16	1.0	16	16384	16K (uniform)	512	8192	1.13	1.28	1.33
Qwen2.5-7B	16	0.5	16	12554	8K–16K	389	6220	1.15	1.39	1.43
Qwen2.5-7B	16	0.1	16	9688	1K–16K	340	5446	1.13	1.55	1.59
Qwen2.5-7B	16	1.0	16	65536	65K (uniform)	512	8192	1.06	1.55	1.52
Qwen2.5-7B	16	0.5	16	45214	32K–65K	379	6060	1.13	1.90	1.87
Qwen2.5-7B	16	0.1	16	38252	6K–65K	332	5314	1.18	2.35	2.21
Qwen2.5-7B	32	1.0	32	16384	16K (uniform)	512	16384	1.14	1.44	1.44
Qwen2.5-7B	32	0.5	32	12052	8K–16K	391	12521	1.26	1.73	1.61
Qwen2.5-7B	32	0.1	32	8796	1K–16K	271	8672	1.32	2.21	2.00
Qwen2.5-7B	32	1.0	32	65536	65K (uniform)	512	16384	1.07	1.66	1.67
Qwen2.5-7B	32	0.5	32	45750	32K–65K	385	12313	1.12	2.20	2.15
Qwen2.5-7B	32	0.1	32	34370	6K–65K	288	9226	1.21	3.22	2.68
Llama-3.1-8B	1	1.0	1	65535	65K (uniform)	512	512	1.24	1.59	1.59
Llama-3.1-8B	16	1.0	16	16383	16K (uniform)	512	8192	1.26	1.48	1.45
Llama-3.1-8B	16	0.5	16	12553	8K–16K	389	6220	1.30	1.63	1.59
Llama-3.1-8B	16	0.1	16	9687	1K–16K	340	5446	1.43	1.92	1.91
Llama-3.1-8B	16	1.0	16	65535	65K (uniform)	512	8192	1.13	1.67	1.68
Llama-3.1-8B	16	0.5	16	45213	32K–65K	379	6060	1.24	2.09	2.09
Llama-3.1-8B	16	0.1	16	38251	6K–65K	332	5314	1.23	2.69	2.47
Llama-3.1-8B	32	1.0	32	16383	16K (uniform)	512	16384	1.30	1.61	1.57
Llama-3.1-8B	32	0.5	32	12051	8K–16K	391	12521	1.46	1.89	1.88
Llama-3.1-8B	32	0.1	32	8795	1K–16K	271	8672	1.52	2.74	2.35
Llama-3.1-8B	32	1.0	32	65535	65K (uniform)	512	16384	1.16	1.93	1.93
Llama-3.1-8B	32	0.5	32	45749	32K–65K	385	12313	1.24	2.31	2.31
Llama-3.1-8B	32	0.1	32	34369	6K–65K	288	9226	1.33	3.62	2.91

Lean wins across every measured cell — all 26 have tps× ≥ 1.06, none regress. The win grows with raggedness (lower ratio) and context: peak 1.52× throughput (Llama b32/16K r0.1) and 3.62× / 2.91× lower median / p99 ITL (Llama b32/65K r0.1). Uniform batches are also faster at every size (1.06–1.30× throughput, 1.28–1.93× lower ITL) — e.g. b1 uniform 65K: Qwen 1.27× tps / 1.63× ITL, Llama 1.24× / 1.59×; b32 uniform 65K: Qwen 1.07× / 1.66×, Llama 1.16× / 1.93×.

Command (per arm; identical --seed 42):

# Server — baseline prepends `SGLANG_DISABLE_LEAN_ATTENTION=1`; Lean is the default (no flag) or `--enable-lean-attention`:
python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct \
  --attention-backend triton --disable-radix-cache --context-length 73728 \
  --mem-fraction-static 0.85 --cuda-graph-bs 1 2 4 8 16 32 --port 30000
# Client — per (B, ctx, ratio):
python3 -m sglang.bench_serving --backend sglang --model Qwen/Qwen2.5-7B-Instruct \
  --dataset-name random --random-input-len <ctx> --random-output-len 512 \
  --random-range-ratio <1.0|0.5|0.1> --num-prompts <B> --max-concurrency <B> --seed 42

End-to-end serving — MLA (DeepSeek-R1-MXFP4, TP=4, fp8 KV)

Same closed-loop methodology, --tp-size 4 --kv-cache-dtype fp8_e4m3. Rows include batch 1 to show the gate boundary and where Lean helps vs regresses. The auto-gate enables MLA at batch ≥ 8 (below it, the standard kernel runs by default).

B	ratio	reqs	KV mean	KV range	dec/req	dec iters	tps×	ITL×	p99×
1	1.0	1	65536	65K (uniform)	512	512	0.83	0.60	0.60
8	1.0	8	65536	65K (uniform)	512	4096	1.06	1.33	1.42
8	0.1	8	41112	6K–65K	306	2448	1.09	1.66	1.98
16	1.0	16	65536	65K (uniform)	512	8192	1.07	1.42	1.66
16	0.1	16	38252	6K–65K	332	5314	1.10	1.87	2.70
32	1.0	32	65536	65K (uniform)	512	16384	1.05	1.36	1.34
32	0.1	32	34370	6K–65K	288	9226	1.13	2.41	2.47

Lean wins at batch ≥ 8 (the gate threshold): 1.05–1.13× tps, 1.33–2.41× lower ITL, growing with raggedness (peak b32 ragged: 1.13× tps, 2.41× ITL). At batch 1 it regresses (0.83× tps, 0.60× ITL), confirming the gate is correct — MLA's splitK is faster at low batch, so the auto-gate keeps it off there.

Command:

python3 -m sglang.launch_server --model-path <deepseek-r1-mxfp4> --tp-size 4 \
  --trust-remote-code --attention-backend triton --kv-cache-dtype fp8_e4m3 \
  --disable-radix-cache --mem-fraction-static 0.8 --context-length 163840 --port 9000
python3 -m sglang.bench_serving --backend sglang --model <deepseek-r1-mxfp4> --dataset-name random \
  --random-input-len <ctx> --random-output-len 512 --random-range-ratio <1.0|0.5|0.1> \
  --num-prompts <B> --max-concurrency <B> --seed 42

Default path (auto-gate, no flag)

The tables above are Lean-vs-baseline in the region the auto-gate enables. Confirming the default path (no --enable-lean-attention) reproduces them under CUDA graphs — 3 arms, Qwen GQA, 64K, ragged ratio 0.1, --seed 42: auto (no flag) tracks forced where the capture policy bakes Lean (b≥16) and stays at the standard kernel below the threshold:

batch	baseline tok/s	auto tok/s	forced tok/s	auto tracks
8	79.55	79.53	88.83	baseline (Lean off, b<16)
16	92.67	103.76	107.53	forced (Lean on)
32	96.28	112.77	115.07	forced (Lean on)

So the throughput/latency gains are delivered by default, with no configuration change. (At b8 the forced arm is already faster than baseline — 88.83 vs 79.55, ITL 1.98× — so the b≥16 GQA capture threshold is deliberately conservative and leaves some b8 headroom on the table; tightening it is future work.)

MI300X (gfx942, 304 CUs)

All tests: 1× AMD Instinct MI300X (304 CUs, gfx942, 192 GB HBM3), ROCm 7.2.0, PyTorch2.9.1+rocm7.2.0, Triton 3.7.0 (AMD wheel), CUDA graphs on, --seed 42, identical harnesses, commands, and auto-gate as above. Parity is arch-independent and passes unchanged (bf16 cos 1.0000, fp8 >0.99, paged >0.999). Results are at the default 1×CU grid
(SGLANG_FORCE_LEAN_GRID_CU_MULT default 1.0, one CTA per CU).

Kernel-level (speedup = standard ÷ lean; fp16, page_size=1)

Qwen2.5-7B (28Q/4KV):

batch	8K	16K	32K	64K	128K
1	0.53×	0.66×	1.38×	2.29×	3.72×
8	0.53×	1.03×	1.24×	1.27×	1.25×
32	0.96×	1.14×	1.16×	1.34×	1.37×

Llama-3.1-8B (32Q/8KV):

batch	8K	16K	32K	64K	128K
1	0.53×	0.85×	1.64×	2.59×	3.34×
8	0.96×	1.28×	1.35×	1.38×	1.34×
32	1.20×	1.23×	1.24×	1.29×	—

Peak 3.72× (Qwen) / 3.34× (Llama) at batch=1 / 128K; high batch now wins across context. Short-context low-batch stays below 1× and is kept on the standard kernel by the gate.

End-to-end serving — GQA (auto-gate enabled region, batch ≥ 16)

Same closed-loop methodology and columns as the MI355X table above. KV mean / KV range / dec-per-req / dec iters are identical to the MI355X rows by construction (same --seed 42, model, random-input-len, random-range-ratio); only tps× / ITL× / p99× are re-measured on MI300X.

model	B	ratio	reqs	KV mean	KV range	dec/req	dec iters	tps×	ITL×	p99×
Qwen2.5-7B	16	1.0	16	16384	16K (uniform)	512	8192	0.98	1.04	1.13
Qwen2.5-7B	16	0.5	16	12554	8K–16K	389	6220	1.01	1.10	1.17
Qwen2.5-7B	16	0.1	16	9688	1K–16K	340	5446	1.09	1.20	1.21
Qwen2.5-7B	32	1.0	32	16384	16K (uniform)	512	16384	1.01	1.09	1.08
Qwen2.5-7B	32	0.5	32	12052	8K–16K	391	12521	1.07	1.18	1.21
Qwen2.5-7B	32	0.1	32	8796	1K–16K	271	8672	1.07	1.29	1.28
Qwen2.5-7B	16	1.0	16	65536	65K (uniform)	512	8192	1.02	1.22	1.17
Qwen2.5-7B	16	0.5	16	45214	32K–65K	379	6060	1.04	1.35	1.32
Qwen2.5-7B	16	0.1	16	38252	6K–65K	332	5314	1.04	1.59	1.51
Qwen2.5-7B	32	1.0	32	65536	65K (uniform)	512	16384	1.00	1.14	1.15
Qwen2.5-7B	32	0.5	32	45750	32K–65K	385	12313	1.04	1.37	1.36
Qwen2.5-7B	32	0.1	32	34370	6K–65K	288	9226	1.10	1.99	1.67
Llama-3.1-8B	16	1.0	16	16383	16K (uniform)	512	8192	1.01	1.03	1.02
Llama-3.1-8B	16	0.5	16	12553	8K–16K	389	6220	1.06	1.11	1.11
Llama-3.1-8B	16	0.1	16	9687	1K–16K	340	5446	1.12	1.27	1.27
Llama-3.1-8B	32	1.0	32	16383	16K (uniform)	512	16384	1.05	1.13	1.09
Llama-3.1-8B	32	0.5	32	12051	8K–16K	391	12521	1.11	1.25	1.23
Llama-3.1-8B	32	0.1	32	8795	1K–16K	271	8672	1.13	1.46	1.41
Llama-3.1-8B	16	1.0	16	65535	65K (uniform)	512	8192	1.08	1.17	2.81
Llama-3.1-8B	16	0.5	16	45213	32K–65K	379	6060	1.12	1.38	1.36
Llama-3.1-8B	16	0.1	16	38251	6K–65K	332	5314	1.14	1.66	1.52
Llama-3.1-8B	32	1.0	32	65535	65K (uniform)	512	16384	1.11	1.77	1.71
Llama-3.1-8B	32	0.5	32	45749	32K–65K	385	12313	1.12	1.65	1.32
Llama-3.1-8B	32	0.1	32	34369	6K–65K	288	9226	1.14	2.14	1.83

At 1×CU the enabled region wins across the board on MI300X — including uniform batches (neutral-to-positive throughput, better tail latency). Peak 1.14× throughput and 2.14× / 1.83× lower median / p99 ITL (Llama b32×64K ragged). The MI355X auto-gate thresholds transfer unchanged.

End-to-end serving — MLA (proxy: DeepSeek-V2-Lite, fp8 KV)

amd/DeepSeek-R1-MXFP4 cannot run on MI300X: fp4/MXFP4 matmul is a CDNA4 (gfx950) feature — gfx942 has no fp4 GEMM (0 AFP4WFP4 configs; aiter marks the OCP fp4 dtypes "gfx950+"). MLA is therefore proxied on MI300X with DeepSeek-V2-Lite (16B MoE, MLA, fp8 KV, TP=1), same ragged sweep.
⚠️ V2-Lite has 16 query heads vs R1's 128, so its low-batch CU-saturation differs — this transfers to the R1 MLA gate only qualitatively.

B	ratio	ctx	tps×	ITL×	p99×
8	1.0	65K	1.04	1.12	1.12
8	0.5	65K	0.99	1.16	1.12
8	0.1	65K	0.94	1.09	0.90
16	1.0	65K	1.07	1.29	1.28
16	0.5	65K	1.09	1.50	1.48
16	0.1	65K	1.06	1.37	1.70
32	1.0	65K	1.08	1.36	1.36
32	0.5	65K	1.14	1.75	1.74
32	0.1	65K	1.12	1.67	1.55

b1 loses (0.31–0.42× — Lean's fixed grid is pure overhead at V2-Lite's 16-head batch-1 decode); b8 is borderline (~parity to 1.04×); b16/b32 win (up to 1.14× tps, 1.75× ITL). The clean-win region on this proxy is b ≥ 16 (b8 borderline). Note this differs from R1 on MI355X, where the 128-head decode makes batch 1 a win at 1×CU — so the exact is_mla batch threshold is a function of head count and should be set from the R1 calibration, not this proxy. (b1 rows omitted from the table.)

Conclusion

Lean Attention is a flat-latency, work-stealing decode kernel for long-context and ragged serving on ROCm — up to ~5.9× at the kernel level, 1.52× end-to-end throughput and up to 3.6× lower median / 2.9× lower p99 ITL on GQA, and 1.13× tps / 2.41× lower ITL on MLA at batch ≥ 8 — at zero accuracy cost (bit-identical bf16; fp8/paged parity within quantization noise). It supports MHA/GQA/MLA, bf16/fp16/fp8 KV, and contiguous or paged KV, and runs behind a conservative, CUDA-graph-aware auto-gate that turns it on only where it wins (delivering the gains on the default path) and leaves the standard kernel in place everywhere else — so there is no regression risk.

Reproduce: kernel sweep — SGLANG_FORCE_LEAN_GRID_CU_MULT=1.0 python3 benchmark/lean_kernel_sweep.py;
CI parity (bf16/fp8/paged) — pytest test/registered/kernels/test_lean_attention.py; end-to-end — the sglang.launch_server + sglang.bench_serving commands in each subsection above.

---

CI States

Latest PR Test (Base): ✅ Run #32810803823
Latest PR Test (Extra): ❌ Run #32810803658
Latest PR Test (AMD ROCm 7.2): ❌ Run #32810803793

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #33576

Merge verdict: ⚠️ Do NOT merge yet — not because of a red X, but because the PR's core code is currently unverified by CI. The Lean decode kernel's only correctness test (test_lean_attention.py) never ran — its job aborted on an earlier unrelated failure before reaching it. Separately, this PR modifies the default TritonAttnBackend decode path (auto-gate + always-allocated buffers), so "green elsewhere" does not prove the standard path is unaffected. AMD PR CI is also incomplete: wait-for-stage-b-amd failed and 5 stage-c AMD jobs were skipped.

Caution

The changed code is not exercised by any PR-CI test that actually ran. test/registered/kernels/test_lean_attention.py is registered only to the stage-b-test-1-gpu-small-amd-mi35x suite. That job failed on test_verify_shared_kv.py and aborted (exit 255) before running the lean test (confirmed: the file was in the planned list at est_time=60 but produced no result). The auto-gate in triton_backend.py:forward_decode also alters the default decode path (enable_lean_attention=None → lean_decode_seqlen_gate / lean_capture_policy), and lean buffers are unconditionally allocated in init_forward_metadata / init_cuda_graph_state. Before merge, re-run the mi35x suite so test_lean_attention.py executes, and confirm the auto-gate's default-path impact.

Caution

AMD PR CI is incomplete. wait-for-stage-b-amd failed, so 5 stage-c AMD jobs (stage-c-test-4-gpu-amd, stage-c-test-large-8-gpu-amd, stage-c-*-mi35x, …) were skipped / not tested. NVIDIA base PR Test did not appear among triggered runs for this head. XPU and Xeon jobs were cancelled. Author should ensure required vendor pipelines finish (rerun / update branch).

Changed files: decode_attention.py (+832), triton_backend.py (+119), environ.py (+13), server_args.py (+5), benchmark/lean_kernel_sweep.py (+163, new), test/registered/kernels/test_lean_attention.py (+453, new)

Executed CI failure attribution: AMD: 5 failures (0 clearly PR-caused, 2 possibly) · NPU: 1 failure (0 related) · Others: gate/finish cascades only (no executed test failures)

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
mi35x-1gpu-small	test/registered/attention/test_verify_shared_kv.py	test_rejects_multiple_local_kv_heads	AssertionError: True is not false	🟡	In the modified file triton_backend.py's import surface, but PR does not touch verify_shared_kv_fwd's rejection logic → likely pre-existing. Critically, this failure blocks the PR's own lean test from running.
mi300-1gpu-small (11)	test/registered/attention/test_deterministic.py	test_prefix_with_logprobs	AssertionError (assert result == 1)	🟡	Exercises the triton backend (attention_backend='triton', enable_lean_attention=None) where the auto-gate is live; lean's cross-CTA atomic combine is a classic determinism risk. Worth a human check.
mi300-1gpu-large	test/registered/perf/test_bench_serving_1gpu_part2.py	test_score_api_latency_throughput	AssertionError: 55.4 not less than 48	🟢	Perf threshold on aiter backend (Qwen3-Reranker), not triton/lean. Latency flake.
mi300-1gpu-small (0)	test/registered/moe/test_torch_compile_moe.py	setUpClass	TypeError: launcher() missing 1 required positional argument: '_grid_2'	🟢	torch-compile/Triton launcher API mismatch at server startup (MoE); environmental — same _dynamo error also appears in the bench job. Not attention.
8-gpu-mi35x-disagg	test/registered/amd/disaggregation/test_disaggregation_basic.py	(server)	RuntimeError: Internal Server Error	🟢	PD-disaggregation server error; not the decode-kernel path.

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
base-c-test-perf-8-npu-a3 (+ 4-npu, 16-npu)	test/registered/npu/performance/minimax_m2_5/test_npu_minimax_m2_5_w8a8_4p_in64k_out1k_prefix90_50ms.py	test_npu_..._50ms	ConnectionRefused → assert_metrics: float() argument ... not 'NoneType'	🟢	NPU (Ascend) server failed to start (port 20266 refused). Different backend entirely; PR is AMD-Triton-only.

PR Test Extra (run) and PR Test Extra (AMD) (run) show only call-gate / pr-gate + *-finish failures — aggregator cascades, no executed test-level failures.

Details / what to do before merge

• Verify the PR's own code runs (highest priority). Re-run stage-b-test-1-gpu-small-amd-mi35x so test_lean_attention.py actually executes. It is the sole coverage for the entire PR; right now it has produced zero signal.
• Check the default-path impact. Because enable_lean_attention=None routes through lean_decode_seqlen_gate/lean_capture_policy, the standard Triton decode path is modified for every user, and lean buffers are always allocated. Confirm this doesn't regress short-context decode or CUDA-graph memory. Consider whether the test_deterministic 🟡 failure reproduces with the triton backend + lean gate.
• Triage the two 🟡 rows (test_verify_shared_kv, test_deterministic) to confirm pre-existing vs PR-induced — a clean baseline was unavailable (recent same-workflow runs on other heads had these AMD jobs skipped/cancelled by fast-fail).
• Ensure vendor CI completeness. Push a fresh run so wait-for-stage-b-amd, stage-c AMD, XPU, and Xeon complete; the 🟢 NPU/perf/torch-compile failures appear environmental but should re-run green.

---

Generated by amd-bot using Claude Code CLI