[TRITON][GLUON][GFX950][DSV4] Paged Sparse Attention Gluon Kernel

[cagrikymk]

gfx950 Paged Sparse Attention for DSV4

fp8 flavors: OCP = float8_e4m3fn (bias 7, native gfx950 cvt); fnuz = float8_e4m3fnuz (bias 8, no native gfx950 cvt).

Gluon supports both on the uniform pool, fnuz is dequanted via bf16 (float8e4b8 → bf16 → f32, exact and cheap on gfx950; the direct fnuz→f32 path lowers to a software unpack that spills). The triton fallback dequant is dtype-driven.

Supported formats & dtypes

pa_decode_sparse picks the backend from arch + inputs: gfx950 with block_h=None → the gluon kernel; otherwise → the shared fallback (the triton kernel, or the gfx1250 gluon kernel on gfx1250). q is bf16/fp16; output is bf16.

Since VLLM uses separated SWA and topk storage, the new gluon supports this format as well.

Also, there are 2 different fp8 storage format supported. The uniform path that stores everything in selected fp8 with fp32 scales. While packed format stores rope in bf16, rest in fp8 and has uint8 scales.

• Two-loop (SWA main + top-k extra_*, the vLLM DSv4 layout) is gfx950-gluon + packed only; everything else is single-loop (the gfx950 uniform pool ignores extra_*; the fallback paths assert it None).
• Only gfx950 gluon handles the 3D packed layouts and both fp8 flavors. The fallback paths (triton / gfx1250 gluon) cover only the 2D uniform pool and assert fnuz — OCP fp8 is rejected there (the gfx1250 kernel's dequant is dtype-driven and could take OCP, but the launcher gate blocks it today).
• fp8 uses 1×64 block scaling: uniform pools carry a separate [pages, D/64] fp32 kv_scales; packed fp8_ds_mla embeds a UE8M0 exponent per 64 NoPE channels. DSv4: D=512 (NoPE 448 + RoPE 64), block=256.

In detail:

Uniform fp8 pool - 2D

• Used by ATOM, uses unified KV cache (for SWA + topk).
• Tensors: unified_kv: [pages, D] fp8 (viewed as uint8), one token per slot (page_size = 1). Scales live in a separate side tensor kv_scales: [pages, D//64] fp32.
• Quantization: the whole head is fp8 with 1×64 block scaling — one fp32 scale per 64-element group (D=512 → 8 groups). fp8 flavor is OCP e4m3 (float8_e4m3fn) or fnuz (float8_e4m3fnuz), selected by the caller.
• Dequant: x = fp8_to_f32(byte) · kv_scales[slot, col//64].
• Addressing: the slot indexes the row directly → data byte offset slot·D + col, scale offset slot·(D//64) + col//64.

Packed fp8 "fp8_ds_mla" -3D

• Used by VLLM, which also uses separate kv cache for SWA and topk parts.
• Tensor: cache: [nb, block, 584] uint8 - a single contiguous buffer of fixed-size pages (block tokens/page); the vLLM DSv4 production KV cache. Scales are embedded, no side tensor.
• Per-page layout (block-planar): [block × 576 data bytes] · [block × 8 scale bytes], so page stride cs0 = block·584.
  • data (576 B/token): [0:448) NoPE fp8 (OCP e4m3) + [448:576) RoPE bf16 (64 elems, not quantized).
  • scale (8 B/token, trailing region): 7 UE8M0 exponent bytes (one per 64-dim NoPE group; 448/64 = 7) + 1 pad byte.
• Dequant (NoPE only): x = fp8_to_f32(byte) · exp2(exp_byte − 127) (UE8M0, bias 127); RoPE is read straight as bf16.
• Addressing (slot → block_idx = slot // block, pos = slot % block):
  • NoPE: block_idx·cs0 + pos·576 + [0:448)
  • RoPE (via bf16 ptr): block_idx·(cs0/2) + pos·288 + 224 + [0:64)
  • scale: block_idx·cs0 + block·576 + pos·8 + group

1. Gluon vs aiter Triton fallback (2D uniform pool) on GFX950

DSv4-Pro config: SWA window 128, index_topk 1024. Cases (tokens/query = SWA + top-k = min(ctx/compress_ratio, 1024)):

• c4a — compress-ratio 4: 128+256 (ctx 1024), 128+1024 (ctx 8192, saturates the 1024 budget).
• c128a — compress-ratio 128: 128+8 (ctx 1024), 128+64 (ctx 8192).

case	ctx	tokens	dtype	gluon µs	triton µs	speedup
c4a	1024	128+256	bf16	19.42	25.87	1.33×
c4a	1024	128+256	fp8 OCP	18.51	35.78	1.93×
c4a	1024	128+256	fp8 fnuz	21.82	174.20	7.98×
c4a	8192	128+1024	bf16	35.10	48.91	1.39×
c4a	8192	128+1024	fp8 OCP	30.71	78.54	2.56×
c4a	8192	128+1024	fp8 fnuz	38.34	503.77	13.14×
c128a	1024	128+8	bf16	13.62	18.79	1.38×
c128a	1024	128+8	fp8 OCP	13.17	24.33	1.85×
c128a	1024	128+8	fp8 fnuz	14.69	91.24	6.21×
c128a	8192	128+64	bf16	13.96	19.89	1.43×
c128a	8192	128+64	fp8 OCP	13.40	25.03	1.87×
c128a	8192	128+64	fp8 fnuz	15.18	92.54	6.10×

• bf16: gluon 1.33–1.43× faster.
• fp8 OCP: gluon 1.9–2.6× faster — the honest same-format fp8 comparison (both use the native gfx950 OCP cvt).
• fp8 fnuz: gluon 6.1–13.1× faster — fnuz has no native gfx950 cvt; triton's software fnuz dequant is very slow (even slower than triton bf16), while gluon stays fast by routing fnuz through bf16. Support is added to keep functional parity with the kernel from vllm
• On gfx950, OCP is the fp8 to use: gluon OCP is fastest overall; gluon fnuz costs +12–25% over OCP (the extra fp8→bf16 dequant step); both fp8 encodings beat bf16.

2. Gluon: format × loop × dtype (fp8 = OCP)

This section shows performance difference between different dtypes and storage formats.

{uniform, packed} × {1-loop, 2-loop} × {fp8, bf16}. 1-loop = SWA+top-k merged into one index/cache; 2-loop = separate SWA (main) + top-k (extra) caches. fp8: uniform = OCP e4m3 + separate fp32 scales; packed = fp8_ds_mla (embedded UE8M0 + RoPE bf16, always OCP).

case	ctx	tokens	uni fp8 1L	pk fp8 1L	pk fp8 2L	bf16 1L†	pk bf16 2L
c4a	1024	128+256	18.52	20.53	20.12	19.27	19.69
c4a	8192	128+1024	30.69	37.56	36.70	35.24	35.98
c128a	1024	128+8	13.09	14.42	18.01	14.02	16.98
c128a	8192	128+64	13.23	15.03	18.50	14.67	17.49

• uniform fp8 (1-loop) is the fastest config — 9–18% under packed fp8 (simple whole-head gather vs packed's embedded-scale + split RoPE gather).
• fp8 < bf16 everywhere (halved KV bandwidth).
• 1-loop ≥ 2-loop, gap grows as top-k shrinks: heavy top-k (c4a) is ~neutral (±2%), but tiny top-k (c128a) pays 19–25% for the 2-loop (per-segment-per-split tile fragmentation). Merging SWA+top-k into one index (1-loop) is always at least as fast.

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4382 --add-label <label>

[github-actions]

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

Suggested change

	import math
	import torch
	import triton
	from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import (
	_pa_decode_sparse as gluon_pa_decode_sparse,
	)
	from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as gluon_pa_decode_sparse_reduce,
	)
	from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import (
	_pa_decode_sparse as _pa_decode_sparse_gfx950,
	)
	from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950,
	)
	from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import (
	_pa_decode_sparse as triton_pa_decode_sparse,
	)
	from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as triton_pa_decode_sparse_reduce,
	)
	from aiter.ops.triton.utils._triton import arch_info
	from aiter.ops.triton.utils.device_info import get_num_sms
	from aiter.ops.triton.utils.logger import AiterTritonLogger
	from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import (
	_pa_decode_sparse as gluon_pa_decode_sparse,
	_pa_decode_sparse_reduce as gluon_pa_decode_sparse_reduce,
	)
	DEVICE_ARCH = arch_info.get_arch()
	import math
	import torch
	import triton
	from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import (
	_pa_decode_sparse as _pa_decode_sparse_gfx950,
	)
	from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950,
	)
	from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import (
	_pa_decode_sparse as gluon_pa_decode_sparse,
	)
	from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as gluon_pa_decode_sparse_reduce,
	)
	from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import (
	_pa_decode_sparse as triton_pa_decode_sparse,
	)
	from aiter.ops.triton._triton_kernels.attention.pa_decode_sparse import (
	_pa_decode_sparse_reduce as triton_pa_decode_sparse_reduce,
	)
	from aiter.ops.triton.utils._triton import arch_info
	from aiter.ops.triton.utils.device_info import get_num_sms
	from aiter.ops.triton.utils.logger import AiterTritonLogger
	DEVICE_ARCH = arch_info.get_arch()