[arai713]

Motivation

gpt-oss attention-sink decode (seqlen_q=1) on gfx942 had two problems. First, it was broken: bf16 sliding-window (SWA-128) sink decode crashed when a windowed config fell into the fp8-loader code path and failed an assert (fp8 total chunks must be divisible by THREADS), even though gfx942 has no fp8 K/V support. Second, the D64 decode path ran at a suboptimal num_warps, leaving significant decode performance on the table. AOTriton has no attention-sink kernel, so there is no competing flash path for gpt-oss sink attention on gfx942, the goal is rocKE-only correctness plus a competitive roofline. This PR fixes the SWA-sink decode crash and tunes decode performance.

Technical Details

• Dispatch fix: guard the gfx942 2D fp8-loader assert and fp8_chunks_per_thread behind KV_FP8, matching the 3D builder. bf16 windowed decode no longer enters the fp8 path.
• Perf tune: _select_2d_num_warps now returns num_warps=1 for gfx942 D64 decode (all_decode), leaving prefill at nw4. Decode is memory-bandwidth-bound, so 4 warps over-subscribe and contend for KV bandwidth with too little parallel work to hide it, a single warp per workgroup streams KV without that overhead. The change is gated to all_decode, so prefill, which is compute-bound, is untouched.
• Benchmark: add attention-sink + sliding-window support to the decode benchmark plus a gpt-oss sink shapes file (gpt_oss_sink_shapes.json).

Performance

Both arches: bf16, D64, GQA 64/8, decode Sq=1. Median of 3 runs, DSL kernel, best across the num_sms sweep. Measured via the Python-lowerer path (C++ engine not built on the bench box), so absolute latencies are a floor.

gfx942 (MI300X)

Same-run nw4 to nw1, 500 iterations, run-to-run variance under 0.1%.

Shape	nw4 (shipped)	nw1 (tuned)	Speedup
full_sink B64 kv2048	558 us	295 us	1.89x
full_sink B128 kv2048	1116 us	334 us	3.34x
full_sink B64 kv8192	2204 us	1158 us	1.90x
swa_sink B128 kv2048	92 us	46 us	2.01x

full_sink decode is the primary win (about 1.9x to 3.3x). SWA decode already sits at a ~46 us floor and is unchanged except at B128. full_sink B1 routes to the 3D path and is unchanged. The nw4 to nw1 ratio is same-session and device-consistent.

gfx950 (MI350X)

No kernel or selector change was needed: the gfx950 selector already picks num_warps=1 for D64 decode, the same config the gfx942 change introduces, so gfx950 is already at the tuned optimum. Latency, 500 iterations, run-to-run variance under 2%.

Shape	Latency	Path
full_sink B1 kv2048	44 us	3d
full_sink B64 kv2048	59 us	3d
full_sink B128 kv2048	119 us	2d
full_sink B64 kv8192	238 us	3d
swa_sink B64 kv2048	30 us	2d
swa_sink B128 kv2048	31 us	2d
swa_sink B64 kv8192	31 us	2d

SWA decode is flat (~30 us) across batch and KV length, confirming window-skipping. AOTriton has no sink kernel, so there is no comparison column.

Test Plan

• Codegen regression test (CPU, no GPU) which verifies the bf16 SWA-decode config no longer trips the fp8 assert.
• Selector unit check where _select_2d_num_warps returns 1 for Sq=1 and 4 for prefill (Sq=512/2048), this confirms decode-only gating.
• Correctness: ran the 8 gpt-oss sink shapes through the decode benchmark, comparing output against the fp32 reference within bf16 tolerance.

Test Result

• Regression + selector checks pass.
• Correctness: all shapes match the fp32 reference, max_abs <= 1.2e-4 (bf16 tol 4e-2), full and SWA, on MI300X.

Related

JIRA ID : AICK-1498

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Changes are behind the feature flag HIPKERNELPROVIDER_ENABLE_ROCKE which is defaulted to OFF.

[therock-pr-bot]

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

[therock-pr-bot]

🎉 All checks passed! This PR is ready for review.

[therock-pr-bot]

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

• python -m pip install pre-commit
• pre-commit install
• pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

[AviralGoelAMD]

Ideally, this

if problem.all_decode:
            return 1

change should be also done on the cpp counterpart that is present at platform/cpp/instances/common/attention_unified_selectors.cpp so that both CPP and PY translators translate down to identical IR.

(I am not sure if we are strictly enforcing that policy just yet but given that its only a two line change on python side, we can do it quickly on cpp as well).

Rest of the PR LGTM.

[AviralGoelAMD]

.

[Copilot]

Pull request overview

This PR fixes correctness and performance issues in rocKE’s attention-sink decode path (notably for gfx942 seqlen_q=1) and extends the decode benchmarking harness to cover attention sinks and sliding-window (SWA) shapes.

Changes:

• Fixes a gfx942 bf16 SWA decode crash by guarding the 2D fp8-loader chunk-count assert behind KV_FP8.
• Tunes gfx942 D64 decode warp selection to prefer num_warps=1 for decode-only (seqlen_q==1 / all_decode), preserving prefill behavior.
• Adds attention-sink + sliding-window coverage to the decode benchmarks and introduces gpt-oss sink shape JSONs; adds a regression test to prevent the prior crash.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp	Adjusts gfx942 D64 decode warp selection (Sq=1 -> nw1).
dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py	Adds regression test ensuring bf16 SWA decode no longer trips the fp8-loader assert.
dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py	Guards fp8 chunk-count assert / chunk computation behind KV_FP8.
dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py	Tunes gfx942 2D warp selection to return nw1 for decode-only problems.
dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/gpt_oss_sink_shapes.json	Adds gpt-oss sink + SWA decode benchmark shapes for gfx950.
dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py	Extends gfx950 decode benchmark to plumb sinks + sliding_window into the DSL path and shapes loader.
dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/gpt_oss_sink_shapes.json	Adds gpt-oss sink + SWA decode benchmark shapes for gfx942.
dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py	Extends gfx942 decode benchmark to plumb sinks + sliding_window into the DSL path and shapes loader; skips Triton when sinks are active.

Comments suppressed due to low confidence (2)

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/decode/benchmark_decode_live.py:228

• sliding_window is now a first-class shape field, but the AITER Triton baseline still hard-codes window_size=(-1, -1). For any shape with sliding_window>0 and use_sinks=False, Triton will benchmark the wrong math (full-causal) and the DSL-vs-Triton comparison becomes invalid.

    except ImportError:

dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/decode/benchmark_decode_live.py:215

• sliding_window is now loaded and passed to the DSL path, but the AITER Triton baseline still hard-codes window_size=(-1, -1). For any shape with sliding_window>0 and use_sinks=False, Triton will benchmark the wrong math (full-causal) and the DSL-vs-Triton comparison becomes invalid.

    except ImportError:
        return None

[aosewski]

Read through this end to end. The fp8 guard is right and the C++ mirror matches (all_decode is defined as max_seqlen_q == 1 in attention_unified_entry.cpp#L125, so the two selectors agree). Four comments inline; these six are on lines outside the diff.

1. --dtype / --batch silently drop the new shape fields (gfx950)

benchmark_decode_live.py#L658-L686 rebuilds DecodeShape field by field, and neither override forwards use_sinks or sliding_window. So this:

python -m benchmarks.gfx950.attention.decode.benchmark_decode_live \
    --shapes .../gpt_oss_sink_shapes.json --batch 64

benchmarks plain full-causal attention under rows still labelled full_sink_* / swa_sink_*. A batch sweep is the most likely way these shapes get run.

shapes = [dataclasses.replace(s, dtype=args.dtype) for s in shapes]
shapes = [dataclasses.replace(s, batch=args.batch) for s in shapes]

immune to the next field someone adds. (gfx942 is safe only because it has neither flag.)

2. AOTriton isn't skipped for sinks or windows

gfx942#L403 / gfx950#L389 — use_bias covers softcap, alibi and qq_bias but not the two new fields, so all 8 shapes in gpt_oss_sink_shapes.json get timed against scaled_dot_product_attention(is_causal=False): no sink logit, no window. That number becomes best_speedup_vs_aoTriton and lands in the printed geomean.

use_bias = (
    alibi_slopes is not None
    or data["qq_bias"] is not None
    or softcap
    or shape.use_sinks
    or shape.sliding_window > 0
)

3. FlyDSL has no gate at all

gfx942#L498-L503 / gfx950#L480-L485 — _flydsl_supported checks head_size and dtype only, and FlyDSL is called with causal=True and no sink or window operand, so flydsl_ms and best_speedup_vs_flydsl are wrong math for these shapes. The prefill sibling already gets this right (benchmark_prefill2d_live.py#L190-L193):

if shape.use_sinks:
    return False, "sinks unsupported"
if shape.sliding_window:
    return False, "sliding_window unsupported"

Between this, #2 and the inline comment on _run_triton, the same two shape fields get three different policies across three backends — one hard skip, one stale predicate, one nothing. A single _backend_supports(shape) helper would keep them honest.

4. window_size is still hardcoded

gfx942#L251 / gfx950#L237 — window_size=(-1, -1), while sliding_window is now a first-class field. Only the sinks flag short-circuits, so a sliding_window>0, use_sinks=False shape (newly expressible because of this PR) times full-causal Triton against a windowed DSL kernel. Copilot flagged this earlier; still open.

5. Same unguarded assert survives in the gfx950 2D builder

gfx950/attention_tiled_2d.py#L2378-L2382 has the identical unconditional fp8_total_chunks % THREADS == 0. With this PR, gfx942 2D, gfx942 3D and gfx950 3D all guard it and gfx950 2D is the last one that doesn't — and this PR adds gfx950 sink-decode shapes that drive that builder. It survives today only because the gfx950 selector happens to land on nw≤2; any retune to nw=4 trips it. Same if KV_FP8: wrap closes it.

6. The correctness result isn't reproducible from this PR

The description reports max_abs <= 1.2e-4 vs an fp32 reference, but neither benchmark computes a reference, compares outputs, or reports an error norm — out is allocated, written and discarded (gfx942#L699-L781). For a change that alters kernel geometry that's the one number worth being able to re-run. Either an opt-in --check that runs a small fp32 sink reference on the first shape, or a pointer to the test that produced those numbers, would do it.

---

Minor, take or leave: the results record (L783-L817) mirrors every other DecodeShape field as its own JSON key but not the two new ones, so downstream tooling has to reverse-parse signature to tell a sink run from a plain one. And the gfx950 file got the code without the comments the gfx942 twin got (sinks tensor and the skip at L214) — they're maintained as clones, so the next person to touch one will wonder.

[aosewski]

LGTM.