Skip to content

feat(rocke): fix and tune attention sink decode #9391

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
arai713 merged 4 commits into from
Jul 30, 2026
Merged

feat(rocke): fix and tune attention sink decode #9391

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
62466f3
Guard gfx942 2D fp8-loader assert behind KV_FP8
arai713 Jul 14, 2026
17c4db5
Tune gfx942 D64 decode to num_warps=1 and bench sinks
arai713 Jul 14, 2026
c8baa47
Add gfx950 sink decode benchmark support
arai713 Jul 16, 2026
baab3ab
fix(rocke): fixes for PR review
arai713 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 43 additions & 10 deletions ...kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,14 +58,21 @@ class DecodeShape:
block_size: int
dtype: str
label: str
use_sinks: bool = False
sliding_window: int = 0
Comment thread
arai713 marked this conversation as resolved.
Resolved

@property
def signature(self) -> str:
return (
sig = (
f"b{self.batch}_sq{self.seqlen_q}_sk{self.seqlen_k}"
f"_nhq{self.num_query_heads}_nhk{self.num_kv_heads}"
f"_hd{self.head_size}_bs{self.block_size}_{self.dtype}"
)
if self.use_sinks:
sig += "_sinks"
if self.sliding_window:
sig += f"_sw{self.sliding_window}"
return sig


def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]:
Expand Down Expand Up @@ -100,6 +107,8 @@ def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]:
block_size=int(merged["block_size"]),
dtype=str(merged.get("dtype", "bf16")),
label=str(merged.get("label", f"kv{merged['seqlen_k']}")),
use_sinks=bool(merged.get("use_sinks", False)),
sliding_window=int(merged.get("sliding_window", 0)),
)
shapes.append(shape)
return shapes
Expand Down Expand Up @@ -180,6 +189,11 @@ def _make_inputs(
if use_qq_bias
else None
)
sinks = (
torch.randn(shape.num_query_heads, dtype=dtype, device="cuda") * 0.1
if shape.use_sinks
else None
)

return dict(
q=q,
Expand All @@ -192,6 +206,7 @@ def _make_inputs(
softcap=softcap,
alibi_slopes=alibi_slopes,
qq_bias=qq_bias,
sinks=sinks,
)


Expand All @@ -209,6 +224,7 @@ def _run_triton(

hip_stream = _bench_stream_handle()
out = torch.empty_like(data["q"])
window_size = (shape.sliding_window - 1, 0) if shape.sliding_window else (-1, -1)

def call_once():
tri(
Expand All @@ -222,15 +238,15 @@ def call_once():
max_seqlen_k=shape.seqlen_k,
softmax_scale=data["scale"],
causal=True,
window_size=(-1, -1),
window_size=window_size,
block_table=data["block_table"],
softcap=data["softcap"],
q_descale=None,
k_descale=None,
v_descale=None,
alibi_slopes=data["alibi_slopes"],
qq_bias=data["qq_bias"],
sinks=None,
sinks=data["sinks"],
)

try:
Expand Down Expand Up @@ -277,6 +293,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters
dtype=shape.dtype,
kv_block_size=shape.block_size,
num_sms=num_sms,
use_sinks=shape.use_sinks,
sliding_window=shape.sliding_window,
)
result = dispatch_attention(req)
path = result.spec.path # "2d" or "3d"
Expand All @@ -296,6 +314,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters
softcap=data["softcap"],
use_alibi=data["alibi_slopes"] is not None,
use_qq_bias=data["qq_bias"] is not None,
use_sinks=data["sinks"] is not None,
sliding_window=shape.sliding_window,
num_sms=num_sms,
)

Expand All @@ -311,6 +331,7 @@ def call_once():
softmax_scale=data["scale"],
block_table=data["block_table"],
softcap=data["softcap"],
sinks=data["sinks"],
alibi_slopes=data["alibi_slopes"],
qq_bias=data["qq_bias"],
backend=run_backend,
Expand All @@ -330,9 +351,9 @@ def _run_aoTriton(
"""Time AOTriton flash SDPA for decode. Returns ms or None on failure.

Reconstructs dense [B, H, S_k, D] KV from the paged cache outside the
timed region. Returns None (skipped) when any bias operand is active
(softcap, alibi_slopes, or qq_bias) because
scaled_dot_product_attention does not support them.
timed region. Returns None (skipped) when any operand it cannot express
is active (softcap, alibi_slopes, qq_bias, attention sinks, or a sliding
window) because scaled_dot_product_attention does not support them.
"""
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
Expand Down Expand Up @@ -367,10 +388,16 @@ def _run_aoTriton(
alibi_slopes = data["alibi_slopes"]
softcap = data["softcap"]

# scaled_dot_product_attention does not support alibi, qq_bias or
# softcap
use_bias = alibi_slopes is not None or data["qq_bias"] is not None or softcap
if use_bias:
# scaled_dot_product_attention does not support alibi, qq_bias,
# softcap, attention sinks, or a sliding window.
unsupported = (
alibi_slopes is not None
or data["qq_bias"] is not None
or softcap
or shape.use_sinks
or shape.sliding_window > 0
)
if unsupported:
return None

# Probe eligibility.
Expand Down Expand Up @@ -469,6 +496,10 @@ def _flydsl_supported(shape: DecodeShape) -> tuple[bool, str]:
return False, f"head_size={shape.head_size} (need 64 or 128)"
if shape.dtype not in ("bf16", "fp16"):
return False, f"dtype={shape.dtype} (need bf16 or fp16)"
if shape.use_sinks:
return False, "sinks unsupported"
if shape.sliding_window:
return False, "sliding_window unsupported"
return True, ""


Expand Down Expand Up @@ -763,6 +794,8 @@ def main() -> int:
"head_size": shape.head_size,
"block_size": shape.block_size,
"dtype": shape.dtype,
"use_sinks": shape.use_sinks,
"sliding_window": shape.sliding_window,
"triton_ms": tri_ms,
"aoTriton_ms": aot_ms,
"dsl": {str(sms): dsl_results[sms] for sms in args.num_sms_sweep},
Expand Down
23 changes: 23 additions & 0 deletions ...kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
{
"meta": {
"model": "gpt-oss-120b/20b",
"arch": "gfx942",
"dtype": "bf16",
"num_query_heads": 64,
"num_kv_heads": 8,
"head_size": 64,
"block_size": 16,
"use_sinks": true,
"note": "gpt-oss sink-bias decode (seqlen_q=1). Two mask families: full-causal (sliding_window=0) and sliding-window SWA-128. bf16/D64/64-8 GQA. AOTriton has no sink kernel, so the Triton column is blank by design (DSL-only). kv_len sweep short->long; batch sweep to fill the GPU (decode is bandwidth-bound only once batched)."
},
"shapes": [
{"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b1_kv2048"},
{"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b64_kv2048"},
{"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 0, "label": "full_sink_b128_kv2048"},
{"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 0, "label": "full_sink_b64_kv8192"},
{"batch": 1, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b1_kv2048"},
{"batch": 64, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b64_kv2048"},
{"batch": 128, "seqlen_q": 1, "seqlen_k": 2048, "sliding_window": 128, "label": "swa_sink_b128_kv2048"},
{"batch": 64, "seqlen_q": 1, "seqlen_k": 8192, "sliding_window": 128, "label": "swa_sink_b64_kv8192"}
]
}
85 changes: 46 additions & 39 deletions ...kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
import math
import sys
import traceback
from dataclasses import dataclass
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Dict, List, Optional

Expand All @@ -52,14 +52,21 @@ class DecodeShape:
block_size: int
dtype: str
label: str
use_sinks: bool = False
sliding_window: int = 0
Comment thread
arai713 marked this conversation as resolved.
Resolved

@property
def signature(self) -> str:
return (
sig = (
f"b{self.batch}_sq{self.seqlen_q}_sk{self.seqlen_k}"
f"_nhq{self.num_query_heads}_nhk{self.num_kv_heads}"
f"_hd{self.head_size}_bs{self.block_size}_{self.dtype}"
)
if self.use_sinks:
sig += "_sinks"
if self.sliding_window:
sig += f"_sw{self.sliding_window}"
return sig


def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]:
Expand Down Expand Up @@ -94,6 +101,8 @@ def load_decode_shapes(paths: List[Path]) -> List[DecodeShape]:
block_size=int(merged["block_size"]),
dtype=str(merged.get("dtype", "bf16")),
label=str(merged.get("label", f"kv{merged['seqlen_k']}")),
use_sinks=bool(merged.get("use_sinks", False)),
sliding_window=int(merged.get("sliding_window", 0)),
)
shapes.append(shape)
return shapes
Expand Down Expand Up @@ -174,6 +183,11 @@ def _make_inputs(
if use_qq_bias
else None
)
sinks = (
torch.randn(shape.num_query_heads, dtype=dtype, device="cuda") * 0.1
if shape.use_sinks
else None
)

return dict(
q=q,
Expand All @@ -186,6 +200,7 @@ def _make_inputs(
softcap=softcap,
alibi_slopes=alibi_slopes,
qq_bias=qq_bias,
sinks=sinks,
)


Expand All @@ -203,6 +218,7 @@ def _run_triton(

hip_stream = _bench_stream_handle()
out = torch.empty_like(data["q"])
window_size = (shape.sliding_window - 1, 0) if shape.sliding_window else (-1, -1)

def call_once():
tri(
Expand All @@ -216,15 +232,15 @@ def call_once():
max_seqlen_k=shape.seqlen_k,
softmax_scale=data["scale"],
causal=True,
window_size=(-1, -1),
window_size=window_size,
block_table=data["block_table"],
softcap=data["softcap"],
q_descale=None,
k_descale=None,
v_descale=None,
alibi_slopes=data["alibi_slopes"],
qq_bias=data["qq_bias"],
sinks=None,
sinks=data["sinks"],
)

try:
Expand Down Expand Up @@ -271,6 +287,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters
dtype=shape.dtype,
kv_block_size=shape.block_size,
num_sms=num_sms,
use_sinks=shape.use_sinks,
sliding_window=shape.sliding_window,
)
result = dispatch_attention(req)
path = result.spec.path # "2d" or "3d"
Expand All @@ -290,6 +308,8 @@ def _run_dsl(shape: DecodeShape, data: dict, num_sms: int, *, warmup: int, iters
softcap=data["softcap"],
use_alibi=data["alibi_slopes"] is not None,
use_qq_bias=data["qq_bias"] is not None,
use_sinks=data["sinks"] is not None,
sliding_window=shape.sliding_window,
num_sms=num_sms,
)

Expand All @@ -305,6 +325,7 @@ def call_once():
softmax_scale=data["scale"],
block_table=data["block_table"],
softcap=data["softcap"],
sinks=data["sinks"],
alibi_slopes=data["alibi_slopes"],
qq_bias=data["qq_bias"],
backend=run_backend,
Expand All @@ -324,9 +345,9 @@ def _run_aoTriton(
"""Time AOTriton flash SDPA for decode. Returns ms or None on failure.

Reconstructs dense [B, H, S_k, D] KV from the paged cache outside the
timed region. Returns None (skipped) when any bias operand is active
(softcap, alibi_slopes, or qq_bias) because
scaled_dot_product_attention does not support them.
timed region. Returns None (skipped) when any operand it cannot express
is active (softcap, alibi_slopes, qq_bias, attention sinks, or a sliding
window) because scaled_dot_product_attention does not support them.
"""
import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
Expand Down Expand Up @@ -361,10 +382,16 @@ def _run_aoTriton(
alibi_slopes = data["alibi_slopes"]
softcap = data["softcap"]

# scaled_dot_product_attention does not support alibi, qq_bias or
# softcap
use_bias = alibi_slopes is not None or data["qq_bias"] is not None or softcap
if use_bias:
# scaled_dot_product_attention does not support alibi, qq_bias,
# softcap, attention sinks, or a sliding window.
unsupported = (
alibi_slopes is not None
or data["qq_bias"] is not None
or softcap
or shape.use_sinks
or shape.sliding_window > 0
)
if unsupported:
return None

# Probe eligibility.
Expand Down Expand Up @@ -459,6 +486,10 @@ def _flydsl_supported(shape: DecodeShape) -> tuple[bool, str]:
return False, f"head_size={shape.head_size} (need 64 or 128)"
if shape.dtype not in ("bf16", "fp16"):
return False, f"dtype={shape.dtype} (need bf16 or fp16)"
if shape.use_sinks:
return False, "sinks unsupported"
if shape.sliding_window:
return False, "sliding_window unsupported"
return True, ""


Expand Down Expand Up @@ -633,35 +664,9 @@ def main() -> int:

shapes = load_decode_shapes(args.shapes)
if args.dtype is not None:
shapes = [
DecodeShape(
batch=s.batch,
seqlen_q=s.seqlen_q,
seqlen_k=s.seqlen_k,
num_query_heads=s.num_query_heads,
num_kv_heads=s.num_kv_heads,
head_size=s.head_size,
block_size=s.block_size,
dtype=args.dtype,
label=s.label,
)
for s in shapes
]
shapes = [replace(s, dtype=args.dtype) for s in shapes]
if args.batch is not None:
shapes = [
DecodeShape(
batch=args.batch,
seqlen_q=s.seqlen_q,
seqlen_k=s.seqlen_k,
num_query_heads=s.num_query_heads,
num_kv_heads=s.num_kv_heads,
head_size=s.head_size,
block_size=s.block_size,
dtype=s.dtype,
label=s.label,
)
for s in shapes
]
shapes = [replace(s, batch=args.batch) for s in shapes]
if args.limit is not None:
shapes = shapes[: args.limit]

Expand Down Expand Up @@ -794,6 +799,8 @@ def main() -> int:
"head_size": shape.head_size,
"block_size": shape.block_size,
"dtype": shape.dtype,
"use_sinks": shape.use_sinks,
"sliding_window": shape.sliding_window,
"triton_ms": tri_ms,
"aoTriton_ms": aot_ms,
"dsl": {str(sms): dsl_results[sms] for sms in args.num_sms_sweep},
Expand Down
Loading
Loading