[HIP] [JIT] fp8_paged_mqa_logits: hand-written gfx950 decode indexer kernel

[sumin-hong]

Validation status: PR kernel commit 81a860a7f2, rebased onto ROCm/aiter@9aa8a6b9.
The original op-level measurements are retained below. The supplied GLM-5.2-MXFP4 E2E and accuracy report is now included under GLM-5.2-MXFP4 E2E and accuracy (2026-09-08 update), using vLLM 8a728663c and AITER 0.1.19 plus this PR's files.

Motivation

The DeepSeek-V3.2 / GLM-5 "lightning indexer" produces the sparse-attention
selection logits. For each query row m and KV position n:

logits[m, n] = sum_h relu(<Q[m, h, :], K[n, :]>) * kv_scale[n] * weights[m, h]

In the decode path K and its per-token dequant scale are co-packed in a paged
cache and gathered through a block table. This PR adds a hand-written HIP kernel
for that path on gfx950 (CDNA4), alongside the existing Triton/Gluon
deepgemm_fp8_paged_mqa_logits. The prefill half is #5046.

The decode indexer streams the entire K cache for a handful of FLOPs per token,
so it is squarely memory-bound and the grid carries a next_n axis -- every one
of a batch's MTP rows re-reads the same K. That redundancy, and the cost of
routing K through LDS, is what this kernel targets.

Technical Details

New op: aiter/ops/fp8_paged_mqa_logits.py

fp8_paged_mqa_logits uses the same tensor interface as deepgemm_fp8_paged_mqa_logits -- same
tensors (q_fp8, kv_cache_fp8, weights, context_lens, block_tables,
max_model_len), same contract that only the causal window
p <= context_lens[b] - next_n + n is written and the -inf outside it is the
caller's. Key design points:

• No LDS on the K path. K is read from the paged cache straight to registers
  and contracted 32 heads x 32 tokens per mfma_scale_f32_32x32x64_f8f6f4 tile.
• ROWS_PER_BLOCK rows share one K stream. R consecutive next_n rows are
  handled by one block, so the cache is read once per R rows instead of once per
  row; grid is (batch, SplitKV, ceil(next_n / R)). R trades that redundancy
  against occupancy, so the host picks it -- along with ChunkK, num_warps and
  SplitKV -- from the total KV footprint batch * max_model_len. Each is
  overridable.
• kv-scale hoist. kv_scale >= 0 and ReLU is positive-homogeneous, so the
  scale is hoisted out of the head sum and applied once per KV column. This also
  matches the order the torch reference applies it in.
• v_permlane32_swap_b32 head reduce, instead of __shfl_down(x, 32) which
  lowers to ds_bpermute_b32 -- an LDS round-trip plus an lgkmcnt wait.
• KVBlockSize is a compile-time constant, instantiated for 1 and 64. Each is
  tied to one cache layout, the pairing production already uses: 1 reads the plain
  co-packed cache, 64 the shuffle_weight(layout=(16,16)) preshuffled one.
• -fno-honor-nans for the module, so the ReLU is a single v_max_f32.
  Without it LLVM must assume a signalling NaN and emits an IEEE canonicalize
  first -- two VALU per accumulator value, in the hottest loop of the kernel.

The kernel is gfx950-only and fixed at n_heads=32, head_dim=128 -- the shipped
GLM-5-FP8 indexer shape. is_supported(num_heads, head_dim, kv_block_size) gates
on that, so a caller that also serves other shapes can route them to the Triton
kernel rather than trip a TORCH_CHECK. This mirrors how
_should_use_asm_kernel gates the head_size=128-only ASM paged-attention kernel
in aiter/ops/attention.py.

Files added / changed:

• aiter/ops/fp8_paged_mqa_logits.py -- the op and its support gate
• csrc/kernels/fp8_paged_mqa_logits.cu -- kernel and host dispatch
• csrc/include/fp8_paged_mqa_logits.h, csrc/pybind/fp8_paged_mqa_logits_pybind.cu
• csrc/include/rocm_ops.hpp, aiter/jit/optCompilerConfig.json -- module_fp8_paged_mqa_logits
• op_tests/test_fp8_paged_mqa_logits.py -- correctness + perf sweep

Test Plan

op_tests/test_fp8_paged_mqa_logits.py runs Triton and HIP on identical inputs
and grades both against one fp32 torch reference (a port of vLLM's
fp8_paged_mqa_logits_torch). Gates are an exact -inf mask match plus
calc_diff < 1e-3 and checkAllclose; tolerances are not widened.

The sweep is the cartesian product of batch in {1,4,16,64},
next_n in {1,2,4,6}, heads in {32,64}, head_dim=128,
kv_len in {1024, 8192, 32768, 131072}, KVBlockSize in {1,64} and
var_ratio in {0.0, 0.3} -- 512 cases, 256 of which the HIP kernel supports.
Points worth calling out:

• next_n up to 6. The host clamps ROWS_PER_BLOCK to next_n, so a sweep
  stopping at 2 can never reach the R=3 instantiation MTP decode actually runs on.
• Ragged context (var_ratio 0.3 draws each length from +/-30% of kv_len).
  A uniform batch gives every sequence the same tail tile and the same causal
  boundary, so a kernel deriving its bounds from one sequence -- or from
  max_model_len -- would pass.
• Shuffled block pool. Block tables are built from a shuffled pool, so a
  kernel that ignores the table and walks the cache linearly fails rather than
  passing by accident.
• Cases the HIP kernel does not support leave its columns nan rather than
  reporting a wrong-but-fast number, and any case dropped for lack of memory is
  logged by name so a short table cannot read as full coverage.

python3 op_tests/test_fp8_paged_mqa_logits.py

Test Result

All correctness gates pass on gfx950 across the sweep; per-case hip err matches
triton err exactly.

Performance on MI355x/gfx950, heads=32, head_dim=128, ragged context
(var_ratio=0.3), run_perftest on an otherwise idle GPU -- 72 shapes:

B	next_n	ctx len	KVBlockSize	Triton µs	HIP µs	speedup
1	2	32768	1	15.8	5.4	2.90x
1	6	131072	1	65.9	14.3	4.61x
4	2	32768	1	40.5	10.2	3.99x
4	6	131072	1	312.9	40.2	7.77x
16	2	32768	1	166.6	29.2	5.71x
16	6	32768	1	394.8	46.9	8.43x
16	6	131072	1	1452.5	212.7	6.83x
64	2	131072	1	1987.8	594.0	3.35x
64	6	131072	1	5122.1	982.8	5.21x
1	2	32768	64	3.9	4.3	0.90x
4	6	131072	64	33.5	27.6	1.21x
16	6	32768	64	39.8	32.7	1.22x
16	6	131072	64	171.4	141.2	1.21x
64	6	8192	64	42.2	32.5	1.30x
64	6	131072	64	712.5	551.7	1.29x

Summarised over the full 72-shape sweep:

group	cases	geomean	range	HIP faster
KVBlockSize=1	36	3.86x	2.01x - 8.43x	36 / 36
KVBlockSize=64 (preshuffled)	36	1.01x	0.69x - 1.61x	18 / 36
all	72	1.97x	0.69x - 8.43x	54 / 72

The win is concentrated on the KVBlockSize=1 path, where the HIP kernel is
faster on every shape measured and the gap widens with next_n (geomean 1.60x at
next_n=1, 1.98x at 2, 2.44x at 6) -- which is what the shared-K-stream design
predicts, since R only has rows to amortise over once next_n > 1.

On the preshuffled KVBlockSize=64 path the two are at parity overall
(geomean 1.01x): the HIP kernel wins by ~1.2-1.3x at high next_n and loses by up
to 0.69x on the small end, where its prologue is not amortised. It is reported
here rather than hidden -- with is_supported() in place a caller can pick per
configuration, and the small-next_n preshuffled corner is the obvious next
target.

Related Work

• [FlyDSL] Paged mla indexer #4221 targets the same paged decode operator with a FlyDSL implementation, broader architecture coverage, and additional KV layouts.

This PR proposes a structurally different, fixed-shape HIP backend with direct register streaming and shared K reads across MTP rows. A direct head-to-head comparison with #4221 and quantitative roofline utilization remain pending.

Current Validation

• Rebased onto upstream main at 9aa8a6b91f1972952314bd16176a76d392f6b85c with no overlapping-file conflicts.
• black==26.3.0 --check: pass for the added Python op and test.
• ruff==0.16.0 check: pass for the added Python op and test.
• python3 -m py_compile: pass for the added Python op and test.
• The supplied integration report below includes a reported op-suite run and model E2E/accuracy results on its stated stack. This documentation update did not rerun GPU experiments; direct comparison with [FlyDSL] Paged mla indexer #4221 and quantitative roofline analysis remain pending.

AI assistance: OpenAI Codex was used for rebase adaptation, formatting, static checks, overlap research, and PR drafting. The submitter reviewed the resulting commit.

GLM-5.2-MXFP4 E2E and accuracy (2026-09-08 update)

This supplied experiment report addresses #5047 (comment). All supplied result tables are retained. The interpretations below distinguish observed differences, unverified statistical estimates, and the accuracy coverage still to be identified.

E2E and accuracy for this PR on GLM-5.2-MXFP4 / MI355X (gfx950) / TP4, per @nholmber's request.

Reported result: about 1.02x output throughput at 128k input on the flat KV layout (block_size=1) across all five concurrency points, with TPOT about 1.03-1.07x better. On the preshuffled layout (block_size=64), throughput is about 0.6% lower across the five points. At 1k input, the reported differences are small and the call counter shows much less decode-indexer activity. These observations describe this MTP=0 integration run.

Setup

model	amd/GLM-5.2-MXFP4, kv-cache-dtype=fp8_e4m3, MTP=0
vLLM	built from source at upstream main 8a728663c (0.28.1rc1.dev388)
aiter	amd-aiter 0.1.19 + this PR's files, JIT module module_fp8_paged_mqa_logits
serve	AMD's GLM-5.2 MXFP4 doc command, --tensor-parallel-size 4, --linear-backend aiter --moe-backend aiter, --block-size {1,64}
bench	vLLM benchmark_serving.py --dataset-name random, --random-range-ratio 1.0, --ignore-eos, --request-rate inf, --seed 1234

Integration — one env-gated branch in rocm_fp8_paged_mqa_logits()
(vllm/v1/attention/ops/rocm_aiter_mla_sparse.py), where vLLM already calls
deepgemm_fp8_paged_mqa_logits. GLM-5.2 is index_n_heads=32, index_head_dim=128, so
is_supported() holds and the kernel really engages.

One caller-side change is required and worth flagging, because a naive drop-in is wrong rather
than slow: this kernel writes only the causal window and leaves the -inf outside it to the
caller. vLLM's existing path relies on a trailing nan_to_num_, which does not clear stale finite values — with this
kernel, stale finite logits from an earlier decode step survive in the reused workspace and corrupt
top-k selection. We pre-fill -inf before the call instead, matching what this PR's own op test does
for both kernels. Cost is symmetric or better: the fill writes 16.9 MB where nan_to_num_ reads and
writes 33.8 MB.

Both arms ran concurrently on one node, one 4-GPU half each, from one container and one aiter
tree. The compared decode paths include the caller-side output initialization described above:

arm	decode paged indexer logits kernel
before	_gluon_deepgemm_fp8_paged_mqa_logits[_preshuffle] (the Gluon reference kernel)
after	aiter.ops.fp8_paged_mqa_logits (this PR)

The branch raises rather than silently falling back when is_supported() is false, and a call
counter prints per TP rank, providing evidence that the HIP path was exercised during the instrumented run:

before : [PR5047] hip calls=0     triton calls=16000
after  : [PR5047] paged MQA logits: heads=32 head_dim=128 block_size=64 -> HIP kernel (aiter#5047)
         [PR5047] hip calls=16000 triton calls=0

op_tests/test_fp8_paged_mqa_logits.py: 512 / 512 configurations reported passing. First JIT build 21.8 s. As described in the original Test Plan, the HIP kernel supports 256 of those configurations; unsupported HIP entries are not counted as HIP coverage.

128k in / 1k out, block_size=1

conc	out tok/s before	after	speedup	mean TPOT ms before	after	TPOT ratio
4	72.2	73.6	1.019x	31.19	29.92	1.042x
8	81.8	83.6	1.022x	71.34	66.92	1.066x
16	90.1	92.5	1.028x	152.71	145.83	1.047x
32	90.5	92.9	1.026x	234.76	228.86	1.026x
64	91.5	93.9	1.025x	235.63	229.57	1.026x

The supplied report summarizes these points as mean +2.40%, with a reported 95% CI of [+2.11%, +2.70%]. Every listed point favours this PR on both metrics. The interval calculation and independent repetition count were not included, so this interval should not be read as an established bound on repeat-run variability.

128k in / 1k out, block_size=64 (preshuffled)

conc	out tok/s before	after	speedup	mean TPOT ms before	after	mean TTFT s before
4	73.0	72.2	0.989x	33.43	33.62	20.5
8	84.2	83.8	0.995x	66.67	67.00	27.8
16	93.7	93.3	0.996x	143.02	144.32	27.2
32	94.7	94.3	0.996x	225.13	226.60	99.9
64	95.6	95.1	0.994x	225.84	226.40	450.7

The supplied report summarizes these points as mean -0.59%, with a reported 95% CI of [-0.85%, -0.33%]. All five throughput differences are negative and TPOT moves in the same direction. This is an observed slowdown in this sweep; the available report does not establish its significance against repeat-run variability.

1k in / 1k out

conc	4	8	16	32	64	128	256
out tok/s before (bs=64)	277.0	456.0	857.7	1358.7	2171.9	3354.1	4717.7
out tok/s after (bs=64)	277.5	458.8	854.9	1347.4	2183.5	3375.1	4726.5
speedup (bs=64)	1.00x	1.01x	1.00x	0.99x	1.01x	1.01x	1.00x
speedup (bs=1)	1.01x	1.01x	1.00x	1.00x	0.99x	1.00x	1.00x

This scenario has limited sensitivity to the kernel. The report attributes this to vLLM short-circuiting the decode indexer when
max_seq_len <= index_topk (sparse_attn_indexer_kpool.py:210,717); GLM-5.2 ships
index_topk=2048 and a 1k/1k server caps at 2080 tokens. Our call counter measures 0.19 op calls
per decode step here against >=2.4 at 128k.

The report gives a spread of sigma = 0.507% and a mean difference of +0.21% across these 14 short-context points. They provide a descriptive short-context comparison, but are not a strict A/A control because the measured op-call count is nonzero. Their spread does not independently establish a +-1% detection threshold or rule out GPU-half placement effects at 128k.

Accuracy

Coverage note: the supplied report contains one baseline/HIP accuracy table and does not identify the block size used for these evaluations. Separate model-accuracy coverage for both block_size=1 and 64 has not yet been established from this report.

Validation gates and expectations from the AMD "GLM-5.2 MXFP4 - vLLM status and validation" deck,
run with the deck's own commands, both arms in parallel on one 4-GPU half each.

gate	before	after	expected
coherence (17x23)	PASS, 391	PASS, 391	coherent, clean </think>, 391
GSM8K (1319 q, 5-shot, flexible-extract)	93.03% +- 0.70	93.48% +- 0.68	~92%
GSM8K (strict-match)	93.10% +- 0.70	93.56% +- 0.68	~92%
GPQA-Diamond (198 q, max_tokens=100k, temp 1.0 / top_p 0.95)	91.41% +- 1.99	88.89% +- 2.23	~92%+
RULER niah_single_2 @ 64k (500 samples)	100%	100%	~90%+
RULER niah_single_2 @ 128k (500 samples)	100%	100%	~90%+

GSM8K differs by 0.45 pp; the report gives a combined standard error of 0.98 pp (z ~ 0.46). GPQA decreases by 2.5 pp — 5 questions out of 198 — with a reported combined standard error of 3.0 pp (z ~ 0.84). The report also records 7 of 198 responses reaching the 100k token budget on each arm. Those approximate error comparisons do not demonstrate equivalence, and both GPQA point estimates are below the stated ~92%+ expectation. The individual response outcomes needed for a paired analysis were not included.

RULER is at the ceiling on both arms at both tested lengths. This supports long-context retrieval on these samples. The caller's -inf initialization contract remains important because stale finite logits can alter top-k selection.

Notes

• Why block_size=1 wins and block_size=64 does not. This PR's op-level table predicts both
  signs: geomean 3.86x at KVBlockSize=1 (36/36 shapes) against 1.01x at KVBlockSize=64 (18/36,
  worst 0.69x). MTP=0 pins next_n=1, this kernel's weakest case by construction — the
  ROWS_PER_BLOCK design shares one K stream across R consecutive next_n rows, and with one row
  there is nothing to amortise. The op table shows the same shape: geomean 1.60x at next_n=1
  rising to 2.44x at next_n=6.
• Why E2E and op-level speedups differ. The op-level measurements cover different batches, contexts and next_n values from this MTP=0 serving run. They suggest a mechanism for the flat-layout benefit, but do not quantify the indexer's share of time in this run. In particular, the original report's 0.13% and 2.5% contribution estimates were extrapolations across configurations rather than measurements of the serving run, and are not used to establish an E2E bound here.
• --block-size must be set explicitly. The serving recipe leaves vLLM's default of 16, which
  is_supported() rejects. Set 1 or 64 explicitly to exercise this kernel and use the dispatch counters above to verify the selected path.
• Nominal vs actual concurrency at 128k. KV cache holds 3.23M tokens and each request costs
  ~129k, so about 25 fit. Throughput saturates at ~95 tok/s from conc 16 on and TTFT reaches 450 s at
  conc 64 — those points queue rather than running fully in parallel. Identical on both arms, so the
  comparison holds, but the concurrency label is nominal.
• Prefix caching disabled (--no-enable-prefix-caching), so absolute numbers are not comparable
  to runs that leave it at the default.
• Sample counts. 1k/1k used 10 prompts and 2 warmups per concurrency slot; 128k/1k used 3 and 1
  (block_size=64) and 2 and 1 (block_size=1). The intended comparisons are baseline versus HIP within each block-size/scenario pair; different sample counts limit comparisons across layouts.
• Not compared against [FlyDSL] Paged mla indexer #4221, the FlyDSL implementation of the same operator. This run was scoped
  to before/after against the Gluon reference. We can add it as a third arm if that would help.
• Optional future work (not part of the requested comparison): MTP > 0, which is where this kernel's design actually pays (every
  number here is next_n=1); a block_size=16 instantiation, since 16 is what most vLLM/ROCm
  deployments run and the kernel currently supports only 1 and 64; and the small-next_n preshuffled
  corner, which is where the -0.59% comes from and which this PR already names as the next target.

Submission Checklist

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

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5047 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

[Copilot]

Pull request overview

This PR adds a gfx950-only, fixed-shape (32 heads, head_dim=128) hand-written HIP backend for the decode-phase FP8 paged MQA “indexer logits” op, wired through the aiter JIT/pybind extension path and accompanied by a correctness/perf sweep that A/Bs against the existing Triton/Gluon implementation.

Changes:

• Add fp8_paged_mqa_logits Python op wrapper + support gate for gfx950 / (32,128) / KVBlockSize {1,64}.
• Add HIP kernel + host dispatch + pybind binding and register a new JIT module (module_fp8_paged_mqa_logits) with -fno-honor-nans.
• Add an op_test sweep script comparing Triton vs HIP vs a torch reference across batch/next_n/ctx_len/layout variations.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
aiter/ops/fp8_paged_mqa_logits.py	New Python entry point + is_supported() gate for the gfx950-only decode kernel.
csrc/kernels/fp8_paged_mqa_logits.cu	HIP kernel implementation and Torch-facing dispatch/autotune logic.
csrc/include/fp8_paged_mqa_logits.h	C++ API declaration for the new op.
csrc/pybind/fp8_paged_mqa_logits_pybind.cu	Pybind module entry for the new op.
csrc/include/rocm_ops.hpp	Adds the pybind macro wiring (FP8_PAGED_MQA_LOGITS_PYBIND).
aiter/jit/optCompilerConfig.json	Registers module_fp8_paged_mqa_logits build config and HIP flags.
op_tests/test_fp8_paged_mqa_logits.py	New correctness + perf sweep comparing Triton vs HIP vs torch reference.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[stefanskiasan]

Results from a production-style deployment, in case they help: GLM-5.3 (full model, GlmMoeDsaForCausalLM, 78 layers, own Quark checkpoint: MXFP4 experts incl. shared experts, FP8 block-scaled attention) on 4× MI355X (gfx950), TP4, ROCm 7.2.3, vLLM glm-release fork (27 Aug) + ROCM_AITER_MLA_SPARSE, MTP k=3, FP8 KV, breakable CUDA graphs, kernels JIT-built on top of the AITER 0.1.19 wheel (the 0.1.21 wheel's Gluon kernels do not lower with our Triton 3.7.1, unrelated to this PR).

• op_tests/test_fp8_paged_mqa_logits.py: all 512 configurations pass against the torch reference (batch up to 64, next_n 4, avg_kv up to 131 072, block_size 1 and 64). Note for others: run the test from outside the checkout, otherwise aiter resolves to the source tree and the JIT build of module_aiter_core starts there.
• End-to-end in vLLM (dispatched from rocm_fp8_paged_mqa_logits when is_supported() is true, --block-size 64): needle-in-haystack prompts at 30k / 150k / 400k / 900k tokens all correct, perplexity unchanged (1.3188 vs 1.3197 on our 20-text set), tool calls and thinking-mode switching fine. Decode throughput at short context is unchanged (as expected); the long-context decode comparison is still running.
• Block size 16 would matter for vLLM/ROCm users: the preshuffled indexer path (PR vllm#51216) is enabled for any block size that is a multiple of 16, and 16 is what most deployments run, so is_supported() returns false there. I tried simply adding case 16 to LAUNCH_KB (and 16 to SUPPORTED_KV_BLOCK_SIZES): compiles, but the test reports calc_diff = 3.4e-01, so the kernel assumes more than the dispatch shows (probably the 64-token block per K stream). A 16-token instance would let this kernel run without changing the KV page geometry (LMCache objects are tied to it, see below), if that is feasible on the kernel side.
• Prefill half ([HIP] [JIT] fp8_mqa_logits: hand-written gfx950 prefill indexer kernel #5046) measured separately, numbers posted there.

One deployment note: switching an instance from block 16 to block 64 with an existing LMCache L2 returns 100 % cache hits with corrupted KV, because the stored objects follow the page geometry; a block-size change needs an empty cache.

[stefanskiasan]

Follow-up, end-to-end decode rather than the kernel in isolation (full GLM-5.3, TP4 on 4x MI355X, fp8 KV, block_size 64, MTP k=3, bs=1, 400 output tokens, no prefix cache, Triton paged indexer vs. this PR's HIP kernel, two runs where given):

context	Triton (tok/s)	HIP #5047 (tok/s)
120k	147 / 145	146 / 145
300k	149 / 143	142 / 144
600k	156	147

So for single-stream decode at these lengths the paged indexer is not on the critical path for us (the prefill kernel from #5046 is a different story: 600k-token prefill 285 s -> 93 s). Just a data point on where the decode kernel pays off — presumably larger decode batches.

[nholmber]

Could you provide e2e performance (say 1k/1k conc4-256 powers of 2 and 128k/1k conc4-64) and accuracy data (esp. long-context & reasoning like GPQA Diamond + Needle-in-Haystack) comparing performance against the Gluon reference kernel (block_size = 1 and block_size=64 for the preshuffle kernel)?

[sumin-hong]

Could you provide e2e performance (say 1k/1k conc4-256 powers of 2 and 128k/1k conc4-64) and accuracy data (esp. long-context & reasoning like GPQA Diamond + Needle-in-Haystack) comparing performance against the Gluon reference kernel (block_size = 1 and block_size=64 for the preshuffle kernel)?

I have attached the additional experiment results to the PR description.