@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5426311d91

[Kernel][ROCm][Perf] FlyDSL decode-attention kernel for 4-bit TurboQuant KV cache (#47896)

Signed-off-by: root <root@smci355-ccs-aus-m02-09.cs-aus.dcgpu> Co-authored-by: root <root@smci355-ccs-aus-m02-09.cs-aus.dcgpu> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

aditi-amd committed Aug 11, 2026 at 08:45 UTC 5426311d9115c597c561967ec3f277ee6e970420

13 files changed +6396 -53

tests/kernels/turboquant/__init__.py

tests/kernels/turboquant/test_flydsl_turboquant_decode.py new

+261

@@ -0,0 +1,261 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +"""Correctness test for the FlyDSL TurboQuant 4-bit KV decode kernel.

5 +

6 +The kernel is validated against a pure-PyTorch full-precision attention

7 +oracle rather than another kernel: the synthetic KV cache is built from

8 +known centroid indices, K-norms and V-scales/zeros, so we can dequantize

9 +the exact stored values and run fp32 softmax attention as ground truth.

10 +Both sides therefore see identical quantized values; the only difference

11 +is bf16-vs-fp32 accumulation, which is bounded by ``ATOL`` below.

12 +

13 +Skipped unless running on ROCm/gfx950 with FlyDSL importable.

14 +"""

15 +

16 +import pytest

17 +import torch

18 +

19 +from vllm.platforms import current_platform

20 +

21 +# Guard the ROCm-only import: `vllm.platforms.rocm` pulls in amdsmi and other

22 +# ROCm-specific modules that are not importable on non-ROCm platforms, so it

23 +# must not be imported at module load time on CUDA/CPU CI.

24 +if not current_platform.is_rocm():

25 + pytest.skip(

26 + "TurboQuant FlyDSL decode only runs on ROCm gfx950.",

27 + allow_module_level=True,

28 + )

29 +

30 +from vllm.platforms.rocm import on_gfx950 # noqa: E402

31 +

32 +if not on_gfx950():

33 + pytest.skip(

34 + "TurboQuant FlyDSL decode only runs on ROCm gfx950.",

35 + allow_module_level=True,

36 + )

37 +

38 +from vllm.v1.attention.ops.flydsl_turboquant_decode import ( # noqa: E402

39 + flydsl_turboquant_decode_attention,

40 + is_flydsl_available,

41 + is_flydsl_gqa6_available,

42 +)

43 +

44 +if not is_flydsl_available():

45 + pytest.skip("FlyDSL runtime is not available.", allow_module_level=True)

46 +

47 +HEAD_SIZE = 128

48 +N_CENTROIDS = 16

49 +KEY_DATA_BYTES = HEAD_SIZE // 2 # 4-bit packed key nibbles

50 +DATA_BYTES_PER_SLOT = HEAD_SIZE # packed K (64) + packed V (64)

51 +NUM_SOA_FIELDS = 3

52 +SOA_K_NORM, SOA_V_SCALE, SOA_V_ZERO = 0, 1, 2

53 +

54 +# bf16-vs-fp32 accumulation tolerance (FlyDSL test_pa.py default).

55 +ATOL = 5e-3

56 +

57 +

58 +def _build_cache(

59 + num_seqs, num_kv_heads, seq_len, qg, kv_block_size, seed=0xC0FFEE, device="cuda"

60 +):

61 + """Vectorized builder for a SoA-layout 4-bit KV cache + fp32 ground truth.

62 +

63 + Returns ``(centroids, q_bf16, kv_cache_4d, block_table, seq_lens,

64 + K_ref, V_ref)`` where ``K_ref``/``V_ref`` are the dequantized

65 + full-precision keys/values used by the reference attention.

66 + """

67 + g = torch.Generator(device=device).manual_seed(seed)

68 + hk = num_kv_heads

69 + hq = hk * qg

70 +

71 + centroids = (

72 + torch.randn(N_CENTROIDS, generator=g, dtype=torch.float32, device=device) * 0.5

73 + )

74 + q = (

75 + torch.randn(

76 + num_seqs, hq, HEAD_SIZE, generator=g, dtype=torch.float32, device=device

77 + )

78 + * 0.1

79 + )

80 + q_bf16 = q.to(torch.bfloat16)

81 +

82 + blocks_per_seq = (seq_len + kv_block_size - 1) // kv_block_size

83 + max_bps = blocks_per_seq + 4

84 + block_table = torch.zeros(num_seqs, max_bps, dtype=torch.int32, device=device)

85 + block_table[:, :blocks_per_seq] = (

86 + torch.arange(

87 + num_seqs * blocks_per_seq, dtype=torch.int32, device=device

88 + ).reshape(num_seqs, blocks_per_seq)

89 + + 1

90 + )

91 + total_blocks = num_seqs * blocks_per_seq + 8

92 +

93 + slot_size_aligned = (KEY_DATA_BYTES + 2) + (KEY_DATA_BYTES + 4) # 134

94 + if slot_size_aligned % 2:

95 + slot_size_aligned += 1

96 + bytes_per_block = kv_block_size * hk * slot_size_aligned

97 +

98 + k_idx = torch.randint(

99 + 0,

100 + N_CENTROIDS,

101 + (num_seqs, hk, seq_len, HEAD_SIZE),

102 + generator=g,

103 + dtype=torch.uint8,

104 + device=device,

105 + )

106 + v_idx = torch.randint(

107 + 0,

108 + N_CENTROIDS,

109 + (num_seqs, hk, seq_len, HEAD_SIZE),

110 + generator=g,

111 + dtype=torch.uint8,

112 + device=device,

113 + )

114 + knorm = (

115 + torch.rand(num_seqs, hk, seq_len, generator=g, device=device) * 0.5 + 0.5

116 + ).to(torch.float16)

117 + vscale = (

118 + torch.rand(num_seqs, hk, seq_len, generator=g, device=device) * 0.05 + 0.01

119 + ).to(torch.float16)

120 + vzero = (

121 + (torch.rand(num_seqs, hk, seq_len, generator=g, device=device) - 0.5) * 0.1

122 + ).to(torch.float16)

123 +

124 + k_ref = centroids[k_idx.long()].float() * knorm.float().unsqueeze(-1)

125 + v_ref = v_idx.float() * vscale.float().unsqueeze(-1) + vzero.float().unsqueeze(-1)

126 +

127 + k_packed = k_idx[..., 0::2] | (k_idx[..., 1::2] << 4)

128 + v_packed = v_idx[..., 0::2] | (v_idx[..., 1::2] << 4)

129 +

130 + kv_cache = torch.zeros(

131 + total_blocks, bytes_per_block, dtype=torch.uint8, device=device

132 + )

133 +

134 + seq_ax = (

135 + torch.arange(num_seqs, device=device)

136 + .view(num_seqs, 1)

137 + .expand(num_seqs, seq_len)

138 + )

139 + tok_ax = (

140 + torch.arange(seq_len, device=device).view(1, seq_len).expand(num_seqs, seq_len)

141 + )

142 + blk_for_tok = block_table[seq_ax, tok_ax // kv_block_size]

143 + slot_for_tok = tok_ax % kv_block_size

144 +

145 + h_ax = torch.arange(hk, device=device)

146 + base_data = (

147 + slot_for_tok.unsqueeze(1) * hk * DATA_BYTES_PER_SLOT

148 + + h_ax.view(1, hk, 1) * DATA_BYTES_PER_SLOT

149 + )

150 + blk_b = blk_for_tok.unsqueeze(1).expand(num_seqs, hk, seq_len)

151 + dst_base = blk_b * bytes_per_block + base_data

152 + rng_k = torch.arange(KEY_DATA_BYTES, device=device)

153 + rng_v = torch.arange(KEY_DATA_BYTES, device=device) + KEY_DATA_BYTES

154 + k_dst = (dst_base.unsqueeze(-1) + rng_k.view(1, 1, 1, -1)).reshape(-1)

155 + v_dst = (dst_base.unsqueeze(-1) + rng_v.view(1, 1, 1, -1)).reshape(-1)

156 + flat = kv_cache.view(-1)

157 + flat[k_dst] = k_packed.reshape(-1)

158 + flat[v_dst] = v_packed.reshape(-1)

159 +

160 + meta_off = kv_block_size * hk * DATA_BYTES_PER_SLOT

161 + meta_view = kv_cache.view(torch.float16).view(total_blocks, -1)

162 + meta_off_hw = meta_off // 2

163 + common = h_ax.view(

164 + 1, hk, 1

165 + ) * NUM_SOA_FIELDS * kv_block_size + slot_for_tok.unsqueeze(1)

166 + blk_flat = blk_for_tok.unsqueeze(1).expand(num_seqs, hk, seq_len)

167 + knorm_off = meta_off_hw + common + SOA_K_NORM * kv_block_size

168 + vscale_off = meta_off_hw + common + SOA_V_SCALE * kv_block_size

169 + vzero_off = meta_off_hw + common + SOA_V_ZERO * kv_block_size

170 + meta_view[blk_flat.reshape(-1).long(), knorm_off.reshape(-1).long()] = (

171 + knorm.reshape(-1)

172 + )

173 + meta_view[blk_flat.reshape(-1).long(), vscale_off.reshape(-1).long()] = (

174 + vscale.reshape(-1)

175 + )

176 + meta_view[blk_flat.reshape(-1).long(), vzero_off.reshape(-1).long()] = (

177 + vzero.reshape(-1)

178 + )

179 +

180 + seq_lens = torch.full((num_seqs,), seq_len, dtype=torch.int32, device=device)

181 + kv_cache_4d = kv_cache.view(total_blocks, kv_block_size, hk, slot_size_aligned)

182 + return (

183 + centroids,

184 + q_bf16,

185 + kv_cache_4d,

186 + block_table,

187 + seq_lens,

188 + k_ref.cpu(),

189 + v_ref.cpu(),

190 + )

191 +

192 +

193 +def _reference_attention(q_bf16, k_ref, v_ref, seq_lens, scale):

194 + """fp32 softmax attention on dequantized keys/values (ground truth)."""

195 + num_seqs, hq, d = q_bf16.shape

196 + hk = k_ref.shape[1]

197 + qg = hq // hk

198 + q = q_bf16.float().reshape(num_seqs, hk, qg, d)

199 + out = torch.zeros(num_seqs, hq, d, dtype=torch.float32)

200 + for s in range(num_seqs):

201 + for h in range(hk):

202 + ql = int(seq_lens[s].item())

203 + k = k_ref[s, h, :ql]

204 + v = v_ref[s, h, :ql]

205 + scores = (q[s, h] @ k.T) * scale

206 + m = scores.max(dim=-1, keepdim=True).values

207 + e = torch.exp(scores - m)

208 + p = e / e.sum(dim=-1, keepdim=True)

209 + out[s, h * qg : (h + 1) * qg] = p @ v

210 + return out

211 +

212 +

213 +@pytest.mark.parametrize("num_seqs", [1, 4, 16])

214 +@pytest.mark.parametrize("seq_len", [256, 1024, 4096])

215 +@pytest.mark.parametrize(

216 + "num_kv_heads,qg",

217 + [

218 + (8, 8), # Qwen2.5-72B class

219 + (8, 16), # Qwen3-32B class

220 + pytest.param(

221 + 4,

222 + 6, # MiniMax-M2.5 class (GQA-6 sibling kernel)

223 + marks=pytest.mark.skipif(

224 + not is_flydsl_gqa6_available(),

225 + reason="FlyDSL GQA-6 sibling kernel not available",

226 + ),

227 + ),

228 + ],

229 +)

230 +@pytest.mark.parametrize("kv_block_size", [16, 32])

231 +def test_flydsl_matches_reference(num_seqs, seq_len, num_kv_heads, qg, kv_block_size):

232 + """FlyDSL decode must match fp32 attention on dequantized KV."""

233 + centroids, q_bf16, kv_cache, block_table, seq_lens, k_ref, v_ref = _build_cache(

234 + num_seqs, num_kv_heads, seq_len, qg, kv_block_size

235 + )

236 + scale = 1.0 / (HEAD_SIZE**0.5)

237 + identity = torch.eye(HEAD_SIZE, dtype=torch.float32, device="cuda")

238 +

239 + out = flydsl_turboquant_decode_attention(

240 + query=q_bf16,

241 + kv_cache=kv_cache,

242 + block_table=block_table,

243 + seq_lens=seq_lens,

244 + Pi=identity,

245 + centroids=centroids,

246 + scale=scale,

247 + mse_bits=4,

248 + key_packed_size=KEY_DATA_BYTES + 2,

249 + value_quant_bits=4,

250 + value_packed_size=KEY_DATA_BYTES + 4,

251 + key_fp8=False,

252 + norm_correction=False,

253 + PiT=identity.T.contiguous(),

254 + max_seq_len=seq_len,

255 + max_num_kv_splits=32,

256 + sinks=None,

257 + )

258 +

259 + ref = _reference_attention(q_bf16.cpu(), k_ref, v_ref, seq_lens.cpu(), scale)

260 +

261 + torch.testing.assert_close(out.cpu().float(), ref, atol=ATOL, rtol=0.0)

vllm/platforms/rocm.py

+23 -8

@@ -623,23 +623,38 @@ class RocmPlatform(Platform):

623 if selected_backend is not None:

624 try:

625 backend_class = selected_backend.get_class()

626 - invalid_reasons = backend_class.validate_configuration(

626 + sel_invalid_reasons = backend_class.validate_configuration(

627 device_capability=device_capability,

628 **attn_selector_config._asdict(),

629 )

630 except ImportError:

631 - invalid_reasons = ["ImportError"]

632 - if invalid_reasons:

633 - raise ValueError(

634 - f"Selected backend {selected_backend} is not valid for "

635 - f"this configuration. Reason: {invalid_reasons}"

636 - )

637 - else:

631 + sel_invalid_reasons = ["ImportError"]

632 + if not sel_invalid_reasons:

633 logger.info_once(

634 "Using %s backend (selected via --attention-backend).",

635 selected_backend.name,

636 )

637 return selected_backend.get_path()

638 + # Only tolerate the mismatch for turboquant_* KV-cache layers:

639 + # boundary layers keep the native dtype (served by the selected

640 + # backend) while turboquant_* layers need TURBOQUANT, so no single

641 + # --attention-backend can serve every layer. For any other dtype

642 + # the explicit selection is genuinely invalid -> fail loud.

643 + kv_dtype = attn_selector_config.kv_cache_dtype

644 + if not (kv_dtype is not None and str(kv_dtype).startswith("turboquant")):

645 + raise ValueError(

646 + f"Selected backend {selected_backend} is not valid for "

647 + f"this configuration. Reason: {sel_invalid_reasons}"

648 + )

649 + # NOTE: pass a str (not the list) -- info_once hashes its args.

650 + logger.info_once(

651 + "Selected backend %s is incompatible with this turboquant "

652 + "layer (%s); using the auto-selected per-layer backend. "

653 + "Reason: %s",

654 + selected_backend.name,

655 + attn_selector_config.attn_type,

656 + str(sel_invalid_reasons),

657 + )

658

659 # No selected backend or the selected backend is invalid,

660 # so we try finding a valid backend.

vllm/v1/attention/backends/turboquant_attn.py

+349 -45

@@ -16,6 +16,7 @@ Per-head per-position slot layout:

16 For turboquant_k3v4_nc head_dim=256: [100 bytes key | 512 bytes value] = 612

17 """

19 +import contextlib

20 import functools

21 import math

22 from dataclasses import dataclass

@@ -26,6 +27,7 @@ import torch.nn.functional as F

28 from vllm.config import get_current_vllm_config

29 from vllm.config.cache import CacheDType

30 +from vllm.logger import init_logger

31 from vllm.model_executor.layers.quantization.turboquant.centroids import (

32 get_centroids,

33 )

@@ -47,6 +49,15 @@ from vllm.v1.attention.backends.fa_utils import (

49 is_flash_attn_varlen_func_available,

50 )

51 from vllm.v1.attention.backends.utils import split_decodes_and_prefills

52 +

53 +# FlyDSL TurboQuant decode (AMD gfx950). Auto-selected when FlyDSL is available

54 +# for eligible layers (SoA store + FlyDSL decode + SoA-aware continuation);

55 +# non-gfx950 or ineligible layers use the SoA Triton decode.

56 +from vllm.v1.attention.ops.flydsl_turboquant_decode import (

57 + flydsl_turboquant_decode_attention,

58 + is_flydsl_available,

59 + is_flydsl_gqa6_available,

60 +)

61 from vllm.v1.attention.ops.triton_turboquant_decode import (

62 _tq_full_dequant_kv,

63 _use_fp8_e4b15,

@@ -59,6 +70,8 @@ from vllm.v1.worker.workspace import (

70 is_workspace_manager_initialized,

71 )

73 +logger = init_logger(__name__)

74 +

75 _HAS_FLASH_ATTN = is_flash_attn_varlen_func_available()

76 if _HAS_FLASH_ATTN:

77 from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func

@@ -71,6 +84,24 @@ if _HAS_FLASH_ATTN:

84 _CONTINUATION_DECODE_THRESHOLD = 128

87 +def _soa_imports():

88 + """Lazy import of the HIP-free SoA Triton subset (store / dequant / decode).

89 +

90 + Kept lazy so the default (FlyDSL-off) path never imports these modules.

91 + """

92 + from vllm.v1.attention.ops.turboquant_soa.triton_turboquant_decode import (

93 + _tq_full_dequant_kv as soa_dequant,

94 + )

95 + from vllm.v1.attention.ops.turboquant_soa.triton_turboquant_store import (

96 + triton_turboquant_store as soa_store,

97 + )

98 + from vllm.v1.attention.ops.turboquant_soa.triton_turboquant_unified_attention import ( # noqa: E501

99 + triton_turboquant_decode_attention_soa as soa_decode,

100 + )

101 +

102 + return soa_store, soa_dequant, soa_decode

103 +

104 +

105 def _build_hadamard(d: int, device_str: str) -> torch.Tensor:

106 """Orthonormal Hadamard matrix (Sylvester construction), cached per (d, device).

107

@@ -290,6 +321,10 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

321

322 supports_quant_query_input: bool = False

323

324 + # Lazily populated before cudagraph capture (FlyDSL decode path only).

325 + _arange_cache: torch.Tensor

326 + _cu_2: torch.Tensor

327 +

328 def __init__(

329 self,

330 num_heads: int,

@@ -337,6 +372,17 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

372 vllm_config.attention_config.tq_max_kv_splits_for_cuda_graph

373 )

374

375 + # FlyDSL decode state. Auto-enabled on gfx950 when FlyDSL is available.

376 + self.sliding_window = sliding_window

377 + self.sinks = kwargs.get("sinks")

378 + # Cache max_model_len now (config is available at __init__ but NOT

379 + # during CUDA-graph capture when _ensure_on_device is re-entered).

380 + self._max_model_len = vllm_config.model_config.max_model_len

381 + # SoA store is required by the FlyDSL decode/continuation path, so it

382 + # tracks FlyDSL availability (single switch for the whole pipeline).

383 + self._use_flydsl = is_flydsl_available()

384 + self._soa_store = self._use_flydsl

385 +

386 def _flash_attn_varlen(

387 self,

388 q: torch.Tensor,

@@ -382,6 +428,44 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

428 quantizer is symmetric around zero (sign-flipping a coordinate

429 maps it to the mirror centroid with identical distortion).

430 """

431 + if self._soa_store:

432 + # CUDA-graph capture safety for the FlyDSL decode path on ROCm.

433 + # (1) Pre-allocate _arange_cache / _cu_2 BEFORE any capture; lazy

434 + # allocation during graph replay lands in the HIP graph memory

435 + # pool and yields stale addresses (GPU fault / garbage).

436 + # (2) Pre-warm the WorkspaceManager to its max size before capture

437 + # so mid-capture growth cannot invalidate pointers baked into

438 + # already-captured batch sizes.

439 + _max_len = self._max_model_len

440 + _already_ok = (

441 + hasattr(self, "_arange_cache")

442 + and self._arange_cache.device.type == str(device).split(":")[0]

443 + and self._arange_cache.shape[0] >= _max_len + 2

444 + )

445 + if not _already_ok:

446 + self._arange_cache = torch.arange(

447 + 0, _max_len + 2, device=device, dtype=torch.int32

448 + )

449 + if not hasattr(self, "_cu_2") or self._cu_2.device != torch.device(device):

450 + self._cu_2 = torch.zeros(2, device=device, dtype=torch.int32)

451 + if (

452 + is_workspace_manager_initialized()

453 + and not current_workspace_manager().is_locked()

454 + ):

455 + B_max = self._max_capture_batch_size()

456 + D = self.head_size

457 + Hq = self.num_heads

458 + S = self.max_num_kv_splits

459 + _pre_warm_bytes = (

460 + B_max * Hq * (S * (D + 1) + D) * 4 # fp32 mid_o + fp32 lse

461 + + B_max * Hq * D * 2 # query-dtype output (bf16 = 2 B)

462 + + 512 # alignment padding

463 + )

464 + with contextlib.suppress(AssertionError):

465 + current_workspace_manager().get_simultaneous(

466 + ((_pre_warm_bytes,), torch.uint8)

467 + )

468 +

469 if not hasattr(layer, "_tq_cached"):

470 D = self.head_size

471

@@ -402,6 +486,29 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

486 layer._tq_midpoints = (c_sorted[:-1] + c_sorted[1:]) / 2

487 layer._tq_cached = True

488

489 + def _max_capture_batch_size(self) -> int:

490 + """Largest decode batch we might see at runtime (for workspace pre-warm).

491 +

492 + Take max(cudagraph_capture_sizes, scheduler.max_num_seqs): a forward

493 + pass exceeding the largest captured graph size falls back to eager,

494 + but the workspace is locked and that eager path can still hit batch

495 + sizes up to max_num_seqs. Falls back to 1024 if config is unavailable.

496 + """

497 + try:

498 + cfg = get_current_vllm_config()

499 + candidates: list[int] = []

500 + sizes = cfg.compilation_config.cudagraph_capture_sizes

501 + if sizes:

502 + candidates.append(int(max(sizes)))

503 + sched = getattr(cfg, "scheduler_config", None)

504 + if sched is not None and getattr(sched, "max_num_seqs", None):

505 + candidates.append(int(sched.max_num_seqs))

506 + if candidates:

507 + return max(candidates)

508 + except Exception: # noqa: BLE001

509 + pass

510 + return 1024

511 +

512 def do_kv_cache_update(

513 self,

514 layer: torch.nn.Module,

@@ -590,6 +697,27 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

697 layer: Any,

698 ):

699 """Quantize + store via fused Triton kernel."""

700 + if self._soa_store:

701 + # SoA layout (data region + metadata region separated per block),

702 + # required by the FlyDSL decode kernel. Pure-Triton store; the

703 + # cache tensor shape is identical to the default AoS store, only

704 + # the within-block byte convention differs.

705 + soa_store, _, _ = _soa_imports()

706 + soa_store(

707 + key=key,

708 + value=value,

709 + kv_cache=kv_cache,

710 + slot_mapping=slot_mapping,

711 + PiT=layer._tq_PiT,

712 + midpoints=layer._tq_midpoints,

713 + mse_bits=self.tq_config.key_mse_bits,

714 + key_packed_size=self.tq_config.key_packed_size,

715 + value_quant_bits=self.tq_config.effective_value_quant_bits,

716 + key_fp8=self.tq_config.key_fp8,

717 + centroids=layer._tq_centroids,

718 + norm_correction=self.tq_config.norm_correction,

719 + )

720 + return

721 triton_turboquant_store(

722 key,

723 value,

@@ -722,21 +850,55 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

850 # Slice from pre-built arange (no kernel launch)

851 synth_seq_lens = _arange_cache[cached_len + 1 : seq_len + 1]

852 synth_bt = attn_metadata.block_table[i : i + 1].expand(q_len, -1)

725 - out = triton_turboquant_decode_attention(

726 - query=q_seq,

727 - kv_cache=kv_cache,

728 - block_table=synth_bt,

729 - seq_lens=synth_seq_lens,

730 - Pi=Pi,

731 - centroids=centroids,

732 - scale=self.scale,

733 - mse_bits=self.tq_config.key_mse_bits,

734 - key_packed_size=self.tq_config.key_packed_size,

735 - value_quant_bits=(self.tq_config.effective_value_quant_bits),

736 - key_fp8=self.tq_config.key_fp8,

737 - norm_correction=self.tq_config.norm_correction,

738 - PiT=PiT,

739 - )

853 + if self._soa_store:

854 + # The cache was written in SoA layout (always, for FlyDSL),

855 + # so it MUST be read with the SoA-aware decode. The

856 + # default AoS decode reads k_norm/v_scale/v_zero from

857 + # the wrong offsets in a SoA cache -> garbage

858 + # cached-prefix output ->

859 + # accuracy collapse on every multi-turn / prefix-cached

860 + # (APC) request. Continuation stays on the Triton SoA

861 + # path even when FlyDSL is the main decode kernel

862 + # (FlyDSL is decode-batch only).

863 + out = self._dispatch_decode_soa(

864 + query=q_seq,

865 + kv_cache=kv_cache,

866 + block_table=synth_bt,

867 + seq_lens=synth_seq_lens,

868 + Pi=Pi,

869 + centroids=centroids,

870 + scale=self.scale,

871 + mse_bits=self.tq_config.key_mse_bits,

872 + key_packed_size=self.tq_config.key_packed_size,

873 + value_quant_bits=(

874 + self.tq_config.effective_value_quant_bits

875 + ),

876 + value_packed_size=self.tq_config.value_packed_size,

877 + max_seq_len=int(seq_len),

878 + key_fp8=self.tq_config.key_fp8,

879 + norm_correction=self.tq_config.norm_correction,

880 + PiT=PiT,

881 + sinks=self.sinks,

882 + sliding_window=self.sliding_window,

883 + )

884 + else:

885 + out = triton_turboquant_decode_attention(

886 + query=q_seq,

887 + kv_cache=kv_cache,

888 + block_table=synth_bt,

889 + seq_lens=synth_seq_lens,

890 + Pi=Pi,

891 + centroids=centroids,

892 + scale=self.scale,

893 + mse_bits=self.tq_config.key_mse_bits,

894 + key_packed_size=self.tq_config.key_packed_size,

895 + value_quant_bits=(

896 + self.tq_config.effective_value_quant_bits

897 + ),

898 + key_fp8=self.tq_config.key_fp8,

899 + norm_correction=self.tq_config.norm_correction,

900 + PiT=PiT,

901 + )

902 else:

903 # Large continuation: dequant cached K/V and use

904 # flash_attn for better throughput.

@@ -801,36 +963,84 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

963 v_cached = v_buf[:, :, :alloc_len, :]

964

965 grid = (alloc_len, 1 * Hk)

804 - _tq_full_dequant_kv[grid](

805 - kv_cache,

806 - block_table,

807 - centroids,

808 - k_cached,

809 - v_cached,

810 - k_cached.stride(0),

811 - k_cached.stride(1),

812 - k_cached.stride(2),

813 - v_cached.stride(0),

814 - v_cached.stride(1),

815 - v_cached.stride(2),

816 - kv_cache.stride(0),

817 - kv_cache.stride(1),

818 - kv_cache.stride(2),

819 - block_table.stride(0),

820 - HEAD_DIM=D,

821 - BLOCK_SIZE=block_size,

822 - NUM_KV_HEADS=Hk,

823 - MSE_BYTES=mse_bytes,

824 - KPS=self.tq_config.key_packed_size,

825 - VQB=self.tq_config.effective_value_quant_bits,

826 - VAL_DATA_BYTES=val_data_bytes,

827 - MSE_BITS=self.tq_config.key_mse_bits,

828 - KEY_FP8=1 if self.tq_config.key_fp8 else 0,

829 - BLOCK_D=BLOCK_D,

830 - NORM_CORRECTION=1 if self.tq_config.norm_correction else 0,

831 - FP8_E4B15=_use_fp8_e4b15(device.index or 0),

832 - num_warps=4,

833 - )

966 + if self._soa_store:

967 + # SoA-aware dequant: read the data/metadata-separated SoA cache

968 + # written by the SoA store. Constants must match the store side.

969 + _, soa_dequant, _ = _soa_imports()

970 + key_fp8 = self.tq_config.key_fp8

971 + key_data_bytes = D if key_fp8 else mse_bytes

972 + data_bytes_per_slot = key_data_bytes + val_data_bytes

973 + meta_region_offset = block_size * Hk * data_bytes_per_slot

974 + num_soa_fields = 2 if key_fp8 else 3

975 + soa_k_norm = 0

976 + soa_v_scale = 0 if key_fp8 else 1

977 + soa_v_zero = 1 if key_fp8 else 2

978 + kv_cache_u16 = kv_cache.view(torch.uint16)

979 + soa_dequant[grid](

980 + kv_cache,

981 + kv_cache_u16,

982 + block_table,

983 + centroids,

984 + k_cached,

985 + v_cached,

986 + k_cached.stride(0),

987 + k_cached.stride(1),

988 + k_cached.stride(2),

989 + v_cached.stride(0),

990 + v_cached.stride(1),

991 + v_cached.stride(2),

992 + kv_cache.stride(0),

993 + block_table.stride(0),

994 + HEAD_DIM=D,

995 + BLOCK_SIZE=block_size,

996 + NUM_KV_HEADS=Hk,

997 + MSE_BYTES=mse_bytes,

998 + VQB=self.tq_config.effective_value_quant_bits,

999 + VAL_DATA_BYTES=val_data_bytes,

1000 + MSE_BITS=self.tq_config.key_mse_bits,

1001 + KEY_FP8=1 if key_fp8 else 0,

1002 + KEY_DATA_BYTES=key_data_bytes,

1003 + META_REGION_OFFSET=meta_region_offset,

1004 + NUM_SOA_FIELDS=num_soa_fields,

1005 + SOA_K_NORM=soa_k_norm,

1006 + SOA_V_SCALE=soa_v_scale,

1007 + SOA_V_ZERO=soa_v_zero,

1008 + BLOCK_D=BLOCK_D,

1009 + NORM_CORRECTION=1 if self.tq_config.norm_correction else 0,

1010 + FP8_E4B15=_use_fp8_e4b15(device.index or 0),

1011 + num_warps=4,

1012 + )

1013 + else:

1014 + _tq_full_dequant_kv[grid](

1015 + kv_cache,

1016 + block_table,

1017 + centroids,

1018 + k_cached,

1019 + v_cached,

1020 + k_cached.stride(0),

1021 + k_cached.stride(1),

1022 + k_cached.stride(2),

1023 + v_cached.stride(0),

1024 + v_cached.stride(1),

1025 + v_cached.stride(2),

1026 + kv_cache.stride(0),

1027 + kv_cache.stride(1),

1028 + kv_cache.stride(2),

1029 + block_table.stride(0),

1030 + HEAD_DIM=D,

1031 + BLOCK_SIZE=block_size,

1032 + NUM_KV_HEADS=Hk,

1033 + MSE_BYTES=mse_bytes,

1034 + KPS=self.tq_config.key_packed_size,

1035 + VQB=self.tq_config.effective_value_quant_bits,

1036 + VAL_DATA_BYTES=val_data_bytes,

1037 + MSE_BITS=self.tq_config.key_mse_bits,

1038 + KEY_FP8=1 if self.tq_config.key_fp8 else 0,

1039 + BLOCK_D=BLOCK_D,

1040 + NORM_CORRECTION=1 if self.tq_config.norm_correction else 0,

1041 + FP8_E4B15=_use_fp8_e4b15(device.index or 0),

1042 + num_warps=4,

1043 + )

1044

1045 # Inverse-rotate MSE keys back to original space

1046 if not self.tq_config.key_fp8:

@@ -930,6 +1140,82 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

1140 )

1141 )

1142

1143 + if self._use_flydsl:

1144 + # FlyDSL decode (gfx950, MSE-key, HEAD_SIZE=128, GQA in {6, 8, 16}).

1145 + # GQA-6 routes to the MiniMax sibling kernel. Ineligible layers fall

1146 + # back to SoA Triton decode.

1147 + _gqa = self.num_kv_groups

1148 + flydsl_gqa_ok = (_gqa in (8, 16)) or (

1149 + _gqa == 6 and is_flydsl_gqa6_available()

1150 + )

1151 + flydsl_eligible = (

1152 + not self.tq_config.key_fp8

1153 + and self.tq_config.key_mse_bits == 4

1154 + and self.tq_config.effective_value_quant_bits == 4

1155 + and self.head_size == 128

1156 + and flydsl_gqa_ok

1157 + and self.sinks is None

1158 + and not (self.sliding_window and self.sliding_window > 0)

1159 + )

1160 + if flydsl_eligible:

1161 + return flydsl_turboquant_decode_attention(

1162 + query=query,

1163 + kv_cache=kv_cache,

1164 + block_table=attn_metadata.block_table,

1165 + seq_lens=attn_metadata.seq_lens,

1166 + Pi=Pi,

1167 + centroids=centroids,

1168 + scale=self.scale,

1169 + mse_bits=self.tq_config.key_mse_bits,

1170 + key_packed_size=self.tq_config.key_packed_size,

1171 + value_quant_bits=self.tq_config.effective_value_quant_bits,

1172 + value_packed_size=self.tq_config.value_packed_size,

1173 + max_seq_len=attn_metadata.max_seq_len,

1174 + key_fp8=self.tq_config.key_fp8,

1175 + norm_correction=self.tq_config.norm_correction,

1176 + PiT=PiT,

1177 + mid_o_buf=mid_o_buf,

1178 + output_buf=output_buf,

1179 + lse_buf=lse_buf,

1180 + buf_holder=layer,

1181 + max_num_kv_splits=self.max_num_kv_splits,

1182 + sinks=self.sinks,

1183 + )

1184 + logger.warning_once(

1185 + "TurboQuant FlyDSL ineligible (key_fp8=%s mse_bits=%s vqb=%s "

1186 + "head_size=%s num_kv_groups=%s sinks=%s) -> SoA Triton decode",

1187 + self.tq_config.key_fp8,

1188 + self.tq_config.key_mse_bits,

1189 + self.tq_config.effective_value_quant_bits,

1190 + self.head_size,

1191 + self.num_kv_groups,

1192 + self.sinks is not None,

1193 + )

1194 + return self._dispatch_decode_soa(

1195 + query=query,

1196 + kv_cache=kv_cache,

1197 + block_table=attn_metadata.block_table,

1198 + seq_lens=attn_metadata.seq_lens,

1199 + Pi=Pi,

1200 + centroids=centroids,

1201 + scale=self.scale,

1202 + mse_bits=self.tq_config.key_mse_bits,

1203 + key_packed_size=self.tq_config.key_packed_size,

1204 + value_quant_bits=self.tq_config.effective_value_quant_bits,

1205 + value_packed_size=self.tq_config.value_packed_size,

1206 + max_seq_len=attn_metadata.max_seq_len,

1207 + key_fp8=self.tq_config.key_fp8,

1208 + norm_correction=self.tq_config.norm_correction,

1209 + PiT=PiT,

1210 + mid_o_buf=mid_o_buf,

1211 + output_buf=output_buf,

1212 + lse_buf=lse_buf,

1213 + buf_holder=layer,

1214 + max_num_kv_splits=self.max_num_kv_splits,

1215 + sinks=self.sinks,

1216 + sliding_window=self.sliding_window,

1217 + )

1218 +

1219 result = triton_turboquant_decode_attention(

1220 query=query,

1221 kv_cache=kv_cache,

@@ -951,3 +1237,21 @@ class TurboQuantAttentionImpl(AttentionImpl["TurboQuantMetadata"]):

1237 max_num_kv_splits=self.max_num_kv_splits,

1238 )

1239 return result

1240 +

1241 + def _dispatch_decode_soa(self, **kwargs):

1242 + """SoA-aware Triton decode — fallback for FlyDSL-ineligible layers.

1243 +

1244 + The SoA decode launcher accepts a subset of the FlyDSL kwargs; filter to

1245 + the params it accepts and raise if a *meaningful* (non-None) kwarg would

1246 + be silently dropped (so we never mask a real feature gap).

1247 + """

1248 + import inspect

1249 +

1250 + _, _, soa_decode = _soa_imports()

1251 + accepted = set(inspect.signature(soa_decode).parameters)

1252 + dropped = [k for k, v in kwargs.items() if k not in accepted and v is not None]

1253 + if dropped:

1254 + raise NotImplementedError(

1255 + f"SoA decode does not support kwargs {sorted(dropped)}"

1256 + )

1257 + return soa_decode(**{k: v for k, v in kwargs.items() if k in accepted})

vllm/v1/attention/ops/flydsl_kernels/__init__.py new

@@ -0,0 +1,2 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/v1/attention/ops/flydsl_kernels/tq_decode.py new

+823

@@ -0,0 +1,823 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# TurboQuant decode (FlyDSL) — MI355X / gfx950 / CDNA4

4 +#

5 +# Decode-only paged attention over a TurboQuant SoA KV cache.

6 +# Targets Qwen-class models: HEAD_SIZE=128, GQA group=8, MSE_BITS=4 K, VQB=4 V,

7 +# N_CENTROIDS=16, BLOCK_SIZE=16, partitioned mode.

8 +#

9 +# === Design ===

10 +# * 1 wave (64 lanes) per CTA — simpler than 4-warp; trades parallelism for

11 +# correctness/debug-ability. A later revision will fan out to 4 warps.

12 +# * Flash-Attention-2 online softmax: rescale running PV/sum per K-tile,

13 +# no QK score storage across tiles, no P-LDS round-trip (qk_acc layout

14 +# matches PV B-operand directly).

15 +# * Centroids loaded once into LDS at kernel start (LDS-resident LUT).

16 +# * K dequant → LDS [token, head_dim] (natural layout)

17 +# * V dequant → LDS [head_dim, token] (TRANSPOSED — required so PV MFMA's

18 +# A=V operand can be loaded as 4-bf16 contiguous via 1 ds_read_b64).

19 +#

20 +# MFMA layouts used:

21 +# QK: mfma(A=K, B=Q, C=qk_acc) → C[m=token, n=query], 4 fp32/lane

22 +# = 4 tokens for SAME query → softmax via

23 +# xor_shuffle(16),(32) reduces per-row

24 +# PV: mfma(A=V_T, B=P, C=acc_pv) → C[m=head_dim, n=query], 4 fp32/lane

25 +# = 4 contiguous head_dim cols for SAME

26 +# query → 1 buffer_store of 8 bytes

27 +# per chunk per lane

28 +

29 +from __future__ import annotations

30 +

31 +import flydsl.compiler as flyc

32 +import flydsl.expr as fx

33 +from flydsl._mlir import ir

34 +from flydsl._mlir.dialects import scf as _scf

35 +from flydsl.expr import (

36 + arith,

37 + buffer_ops,

38 + const_expr, # noqa: F401 kept for kernel-DSL API parity

39 + gpu,

40 + range_constexpr,

41 + rocdl,

42 + vector,

43 +)

44 +from flydsl.expr.typing import Int32, T # noqa: F401 Int32 via fx.Int32

45 +from flydsl.runtime.device import get_rocm_arch as get_hip_arch

46 +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr

47 +

48 +# === Constants (Qwen-class TQ decode profile) ===============================

49 +HEAD_SIZE = 128

50 +# default; overridable via build_tq_decode_module(kv_block_size=...)

51 +KV_BLOCK_SIZE = 16

52 +# MFMA tile = 16 tokens (do not change without re-deriving MFMA shapes)

53 +TILE_SIZE = 16

54 +N_CENTROIDS = 16

55 +# default for tests; overridable via build_tq_decode_module(query_group_size=...)

56 +QUERY_GROUP_SIZE = 16

57 +WARP_SIZE = 64

58 +NUM_WARPS = 1

59 +BLOCK_THREADS = NUM_WARPS * WARP_SIZE

60 +KV_COMPUTE_BLOCK = 256 # 16 K-tiles × 16 tokens

61 +

62 +# TQ SoA layout

63 +NUM_SOA_FIELDS = 3

64 +SOA_K_NORM = 0

65 +SOA_V_SCALE = 1

66 +SOA_V_ZERO = 2

67 +KEY_DATA_BYTES = HEAD_SIZE // 2

68 +VAL_DATA_BYTES = HEAD_SIZE // 2

69 +DATA_BYTES_PER_SLOT = KEY_DATA_BYTES + VAL_DATA_BYTES

70 +

71 +# MFMA

72 +MFMA_M = MFMA_N = 16

73 +MFMA_K_BF16_QK = 32 # CDNA4 wide-K (mfma_f32_16x16x32_bf16)

74 +MFMA_K_BF16_PV = 16 # PV's K = token tile = 16

75 +QK_K_CHUNKS = HEAD_SIZE // MFMA_K_BF16_QK # 4

76 +PV_N_CHUNKS = HEAD_SIZE // MFMA_N # 8

77 +

78 +# LDS regions

79 +CENTROID_LDS_BYTES = N_CENTROIDS * 4 # 64

80 +Q_LDS_BYTES = QUERY_GROUP_SIZE * HEAD_SIZE * 2 # 4096

81 +KV_TILE_LDS_BYTES = TILE_SIZE * HEAD_SIZE * 2 # 4096

82 +

83 +LOG2E = 1.4426950408889634

84 +NEG_INF_VAL = float("-inf")

85 +

86 +

87 +def _vsplat_mul(vec, scalar):

88 + s = scalar.ir_value() if hasattr(scalar, "ir_value") else scalar

89 + return vec * vector.broadcast(T.f32x4, s)

90 +

91 +

92 +allocator = None

93 +

94 +

95 +def build_tq_decode_module(

96 + num_seqs: int,

97 + num_kv_heads: int,

98 + num_partitions: int,

99 + max_blocks_per_seq: int = 512,

100 + softmax_scale: float | None = None,

101 + query_group_size: int = QUERY_GROUP_SIZE,

102 + kv_block_size: int = KV_BLOCK_SIZE,

103 + use_hw_v_transpose: bool = False,

104 + tile_groups_per_partition: int = 1,

105 +):

106 + """Build a TQ decode kernel module.

107 +

108 + ``query_group_size`` (= num_query_heads // num_kv_heads) selects the

109 + GQA factor. Supported values: 8 and 16. For 8 the MFMA's 16-row

110 + capacity is half-used; lanes 8..15 compute garbage and are gated out

111 + of all global-memory writes via an OOB-offset trick on the buffer

112 + resource. Q load iterates ``QG // 4`` chunks (= 2 for QG=8, 4 for

113 + QG=16) so we never read past the real Q rows in HBM.

114 +

115 + ``kv_block_size`` is the number of tokens per vLLM cache block. Must

116 + be a multiple of TILE_SIZE (=16). Supported values: 16 and 32. With

117 + block_size=32, two MFMA K-tiles fit in each block; the kernel walks

118 + block-table entries every ``kv_block_size // TILE_SIZE`` tiles and

119 + re-uses the entry for sub-tiles within the same block. SoA metadata

120 + region grows linearly with block_size.

121 +

122 + ``use_hw_v_transpose`` (CDNA4 / gfx950 only) replaces the 32-element

123 + scattered ``ds_write_b16`` V-dequant pattern with a row-major

124 + ``V[token][head_dim]`` LDS layout written as 4 wide ``ds_write_b128``

125 + per lane, then read back into the MFMA A operand via the hardware

126 + ``ds_read_tr16_b64`` transpose. This eliminates ~28 ds_write

127 + instructions per lane per K-tile (32 -> 4) and 8 strided ds_read_b64

128 + per PV chunk (replaces them with 1 hw-transpose read each).

129 +

130 + ``tile_groups_per_partition`` (default 1) controls the FA-2 split-K

131 + granularity. With value G each partition processes ``G * 16`` K-tiles

132 + = ``G * KV_COMPUTE_BLOCK`` tokens, with the existing 16-tile loop

133 + body iterated G times. The launcher uses this to bound

134 + ``num_partitions`` (e.g. cap at 32) for long context: at 32K /

135 + block_size=32 / num_partitions=32, G=5 covers 32*5*256 = 40 960

136 + tokens of worst-case context with grid.z=32 instead of 256.

137 + Behavior at G=1 is bit-identical to the single-partition kernel.

138 + """

139 + assert query_group_size in (8, 16), (

140 + f"query_group_size must be 8 or 16; got {query_group_size}"

141 + )

142 + assert kv_block_size in (16, 32), (

143 + f"kv_block_size must be 16 or 32; got {kv_block_size}"

144 + )

145 + assert kv_block_size % TILE_SIZE == 0

146 + assert int(tile_groups_per_partition) >= 1, (

147 + f"tile_groups_per_partition must be >= 1; got {tile_groups_per_partition}"

148 + )

149 + USE_HW_TR = bool(use_hw_v_transpose)

150 + TGPP = int(tile_groups_per_partition)

151 + PARTITION_EXTENT_TOKENS = TGPP * KV_COMPUTE_BLOCK

152 +

153 + QG = int(query_group_size)

154 + QG_LOAD_ITERS = QG // 4 # 2 for QG=8, 4 for QG=16

155 + OOB_OFFSET = 0x7FFFFFF0 # noqa: F841 ~2GB byte offset; > any plausible buffer

156 +

157 + _BS = int(kv_block_size)

158 + _TILES_PER_BLOCK = _BS // TILE_SIZE # 1 for BS=16, 2 for BS=32

159 +

160 + global allocator

161 + arch = get_hip_arch()

162 +

163 + if softmax_scale is None:

164 + softmax_scale = 1.0 / (HEAD_SIZE**0.5)

165 + _qk_scale = float(softmax_scale)

166 +

167 + # --- Strides ---

168 + _Hq = num_kv_heads * QG

169 + _stride_q_seq = _Hq * HEAD_SIZE

170 + _stride_q_head = HEAD_SIZE

171 + _stride_bt_seq = max_blocks_per_seq

172 +

173 + _data_region_bytes = _BS * num_kv_heads * DATA_BYTES_PER_SLOT

174 + _meta_region_bytes = num_kv_heads * NUM_SOA_FIELDS * _BS * 2

175 + _stride_cache_block = _data_region_bytes + _meta_region_bytes

176 + _meta_region_offset = _data_region_bytes

177 +

178 + _stride_out_part = QG * HEAD_SIZE

179 + _stride_out_head = num_partitions * QG * HEAD_SIZE

180 + _stride_out_seq = num_kv_heads * num_partitions * QG * HEAD_SIZE

181 + _stride_es_seq = num_kv_heads * num_partitions * QG

182 + _stride_ml_seq = _stride_es_seq

183 +

184 + # --- LDS layout ---

185 + allocator = SmemAllocator(None, arch=arch, global_sym_name="tq_smem")

186 + centroid_off = 0

187 + allocator.ptr = CENTROID_LDS_BYTES

188 + q_off = allocator.ptr

189 + allocator.ptr += Q_LDS_BYTES

190 + kv_off = allocator.ptr

191 + allocator.ptr += KV_TILE_LDS_BYTES

192 +

193 + @flyc.kernel

194 + def tq_decode_kernel(

195 + out_ptr: fx.Tensor,

196 + exp_sums_ptr: fx.Tensor,

197 + max_logits_ptr: fx.Tensor,

198 + query_ptr: fx.Tensor,

199 + kv_cache_ptr: fx.Tensor,

200 + centroids_ptr: fx.Tensor,

201 + block_tables_ptr: fx.Tensor,

202 + seq_lens_ptr: fx.Tensor,

203 + ):

204 + # ---- IDs ---------------------------------------------------------

205 + tid = gpu.thread_idx.x

206 + seq = gpu.block_idx.x

207 + kv_h = gpu.block_idx.y

208 + part = gpu.block_idx.z

209 + lane = tid # 0..63

210 + mfma_row = lane & fx.Int32(15) # = query_row (B's N) /

211 + # token (A's M for QK)

212 + mfma_col_grp = lane >> fx.Int32(4) # 0..3, K-group dim

213 +

214 + # ---- Buffer resources -------------------------------------------

215 + q_rsrc = buffer_ops.create_buffer_resource(query_ptr, max_size=True)

216 + kv_rsrc = buffer_ops.create_buffer_resource(kv_cache_ptr, max_size=True)

217 + bt_rsrc = buffer_ops.create_buffer_resource(block_tables_ptr, max_size=True)

218 + sl_rsrc = buffer_ops.create_buffer_resource(seq_lens_ptr, max_size=True)

219 + cent_rsrc = buffer_ops.create_buffer_resource(centroids_ptr, max_size=True)

220 + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True)

221 + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True)

222 + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True)

223 +

224 + # ---- LDS pointers -----------------------------------------------

225 + base = allocator.get_base()

226 + cent_lds = SmemPtr(base, centroid_off, T.f32, shape=(N_CENTROIDS,))

227 + q_lds_i32 = SmemPtr(base, q_off, T.i32, shape=(Q_LDS_BYTES // 4,)).get()

228 + q_lds_i64 = SmemPtr(base, q_off, T.i64, shape=(Q_LDS_BYTES // 8,)).get()

229 + kv_lds_i32 = SmemPtr(base, kv_off, T.i32, shape=(KV_TILE_LDS_BYTES // 4,)).get() # noqa: F841

230 + kv_lds_i64 = SmemPtr(base, kv_off, T.i64, shape=(KV_TILE_LDS_BYTES // 8,)).get()

231 + kv_lds_i16 = SmemPtr(base, kv_off, T.i16, shape=(KV_TILE_LDS_BYTES // 2,)).get()

232 +

233 + # ---- Constants ---------------------------------------------------

234 + c_sq = fx.Int32(_stride_q_seq)

235 + c_qh = fx.Int32(_stride_q_head)

236 + c_qg = fx.Int32(QG)

237 + c_bt = fx.Int32(_stride_bt_seq)

238 + c_block = fx.Int32(_stride_cache_block)

239 + c_meta_off = fx.Int32(_meta_region_offset)

240 + c_data_per_slot = fx.Int32(DATA_BYTES_PER_SLOT)

241 + c_kv_heads = fx.Int32(num_kv_heads)

242 + c_keydata = fx.Int32(KEY_DATA_BYTES)

243 + c_w = fx.Int32(WARP_SIZE)

244 +

245 + NEG_INF = arith.constant(NEG_INF_VAL, type=T.f32)

246 + ZERO_F = fx.Float32(0.0)

247 + ONE_F = fx.Float32(1.0)

248 + LOG2E_C = arith.constant(LOG2E, type=T.f32)

249 + QK_SCALE = arith.constant(_qk_scale, type=T.f32)

250 +

251 + # Helper: unwrap fx wrapper → raw ir.Value (used in STEP F)

252 + def _ival(v):

253 + return v.ir_value() if hasattr(v, "ir_value") else v

254 +

255 + # ===== STEP A: Load centroids → LDS (cooperative, race-safe) =====

256 + c_idx_safe = lane & fx.Int32(N_CENTROIDS - 1)

257 + c_val = buffer_ops.buffer_load(cent_rsrc, c_idx_safe, vec_width=1, dtype=T.f32)

258 + cent_lds.store(c_val, [arith.index_cast(T.index, c_idx_safe)])

259 + gpu.barrier()

260 +

261 + # ===== STEP B: Load pre-rotated q_rot → row-major LDS ============

262 + # q_rsrc points to the externally-rotated q_rot tensor. Load 8 bf16

263 + # per lane per iter (vec_width=4 i32) in QG_LOAD_ITERS iterations.

264 + for c in range_constexpr(QG_LOAD_ITERS):

265 + row_chunk = lane + fx.Int32(c * WARP_SIZE) # 0..QG*16-1

266 + row = row_chunk >> fx.Int32(4) # 0..QG-1

267 + col_b = row_chunk & fx.Int32(15) # 0..15

268 + col_elem = col_b * fx.Int32(8) # 0..120 in bf16

269 + q_off_elem = seq * c_sq + (kv_h * c_qg + row) * c_qh + col_elem

270 + q_v = buffer_ops.buffer_load(

271 + q_rsrc,

272 + q_off_elem // fx.Int32(2),

273 + vec_width=4,

274 + dtype=T.i32,

275 + )

276 + q_lds_byte = row * fx.Int32(HEAD_SIZE * 2) + col_elem * fx.Int32(2)

277 + vector.store(

278 + q_v,

279 + q_lds_i32,

280 + [arith.index_cast(T.index, q_lds_byte // fx.Int32(4))],

281 + )

282 + gpu.barrier()

283 +

284 + # ===== STEP C: Pre-load Q operands for QK_K_CHUNKS = 4 K-chunks ===

285 + # Wide-K MFMA: K_per_chunk = 32, K_per_lane = 32/4 = 8 bf16 per chunk.

286 + # Lane t holds 8 bf16 at row=mfma_row, K-cols

287 + # chunk*32 + mfma_col_grp*8 + 0..7

288 + # Byte addr = mfma_row * HEAD_SIZE*2 + (chunk*32 + col_grp*8)*2

289 + # = mfma_row * 256 + chunk*64 + col_grp*16

290 + # i64 idx = mfma_row * 32 + chunk*8 + col_grp*2 (each i64 = 4 bf16)

291 + # Load 16 bytes (= 8 bf16 = 2 i64) per chunk via vec(2, i64).

292 + q_chunks = []

293 + for chk in range_constexpr(QK_K_CHUNKS):

294 + q_idx_i64 = (

295 + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8)

296 + + fx.Int32(chk * 8)

297 + + mfma_col_grp * fx.Int32(2)

298 + )

299 + qv = vector.load_op(

300 + T.vec(2, T.i64),

301 + q_lds_i64,

302 + [arith.index_cast(T.index, q_idx_i64)],

303 + )

304 + q_chunks.append(vector.bitcast(T.vec(8, T.bf16), qv))

305 +

306 + # ===== STEP D: Online softmax + PV state =========================

307 + running_max = NEG_INF

308 + running_sum = ZERO_F

309 + zero_v4 = arith.constant_vector(0.0, T.f32x4)

310 + acc_pv = [zero_v4 for _ in range(PV_N_CHUNKS)]

311 +

312 + # ===== STEP E: Sequence-len + partition base ====================

313 + seq_len = buffer_ops.buffer_load(sl_rsrc, seq, vec_width=1, dtype=T.i32)

314 + # ── Bounded num_partitions + internal looping ───────────────────

315 + # With TGPP > 1 this CTA owns a contiguous slice of length

316 + # ``PARTITION_EXTENT_TOKENS`` that is iterated as ``TGPP`` groups

317 + # of 16 K-tiles each. ``partition_start`` is recomputed per

318 + # outer ``tg`` iteration below; the FA-2 online-softmax state

319 + # (running_max, running_sum, acc_pv) accumulates across all

320 + # ``TGPP * 16`` tiles for this CTA — semantically identical to

321 + # processing one large 16*TGPP-tile partition.

322 + partition_base = part * fx.Int32(PARTITION_EXTENT_TOKENS)

323 +

324 + # Per-K-tile dequant lane assignment:

325 + # 16 tokens × 64 packed bytes per token = 1024 bytes

326 + # 64 lanes × 16 bytes/lane = 1024 ✓

327 + # Lane t: token = t / 4, sub-chunk = t % 4 (0..3) → 16 bytes contiguous

328 + tok_in_tile = lane >> fx.Int32(2)

329 + chunk_in_tok = lane & fx.Int32(3)

330 + # ``slot`` (= absolute slot index within the cache block, range

331 + # 0.._BS-1) is recomputed per tile inside the K-tile loop below

332 + # so it reflects ``tile_in_block * TILE_SIZE + tok_in_tile``.

333 +

334 + c_meta_u16_per_kvh = fx.Int32(NUM_SOA_FIELDS * _BS)

335 + c_knorm_off_u16 = fx.Int32(SOA_K_NORM * _BS)

336 + c_vscale_off_u16 = fx.Int32(SOA_V_SCALE * _BS)

337 + c_vzero_off_u16 = fx.Int32(SOA_V_ZERO * _BS)

338 +

339 + # ===== STEP F: K-tile loop =======================================

340 + # With kv_block_size > TILE_SIZE there are _TILES_PER_BLOCK tiles

341 + # per cache block. Tiles within the same block share a block-table

342 + # entry but address different slot ranges within the block.

343 + #

344 + # ── Per-tile block-table OOB redirect ────────────────────────────

345 + # The K-tile loop is unrolled 16 times so every CTA issues 16

346 + # block-table reads regardless of (partition, seq_len). When a

347 + # partition extends past ``seq_len`` (e.g. num_partitions=2 for

348 + # a 256-token sequence — partition 1 covers tokens 256..511, all

349 + # masked out) the per-tile bt offset can land beyond the bt

350 + # allocation. With ``max_size=True`` the descriptor advertises

351 + # 4 GB so the HW returns whatever pre-existing HBM bytes live at

352 + # that offset; the resulting garbage ``phys_block`` is multiplied

353 + # by ``stride_cache_block`` and the subsequent kv_cache read

354 + # jumps into pages that may be unmapped (→ "Memory access fault

355 + # by GPU node") or may decode as NaN bytes (→ NaN segm_out).

356 + # Allocator-pattern dependent, but reproduces deterministically

357 + # at large B (e.g. B=64 seq=256 in the test sweep produces 8 192

358 + # NaN entries even on canonical QG=8).

359 + #

360 + # Fix: when the tile starts at or past ``seq_len`` (so its

361 + # qk_acc is going to be killed by the per-token

362 + # ``kv_tok < seq_len`` mask anyway), redirect the bt read to

363 + # ``bt[seq, 0]`` — always in bounds. The redundant phys_block

364 + # decode + kv_cache read is wasted work, but correctness is

365 + # preserved without changing the kernel's iteration count.

366 + # ===== Runtime-adaptive scf.ForOp w/ iter_args =========

367 + # Runtime-adaptive outer loop: trip count derives from

368 + # actual seq_len, NOT TGPP_max. At cudagraph-capture warmup

369 + # (seq_len=1) only 1 iteration runs; at 32K production

370 + # decode all TGPP iterations run. Kernel binary stays small

371 + # (single body, looped at runtime) — capture time matches

372 + # the legacy 16-tile-only kernel rather than scaling 4x.

373 + #

374 + # FA-2 state (running_max, running_sum, acc_pv[8]) threads

375 + # through scf iter_args; the body reads them at top, runs

376 + # the unchanged 16-tile inner body, and yields the new

377 + # state at bottom. After the loop, results are pulled out

378 + # of for_op.results back into the local Python names so the

379 + # downstream STEP G (output) is unmodified.

380 + c_kcb = fx.Int32(KV_COMPUTE_BLOCK)

381 + c_tgpp = fx.Int32(TGPP)

382 + c_zero_i32 = fx.Int32(0)

383 + c_one_i32 = fx.Int32(1)

384 + # remaining = seq_len - partition_base (signed, may be <=0)

385 + remaining = seq_len - partition_base

386 + in_range = remaining > c_zero_i32

387 + # trip_raw = ceil(remaining / KV_COMPUTE_BLOCK)

388 + # (when in_range is false, the divisor branch produces

389 + # garbage that the select below overrides with zero, so we

390 + # don't pre-clamp)

391 + trip_raw = (remaining + c_kcb - c_one_i32) // c_kcb

392 + trip_clamped = (trip_raw > c_tgpp).select(c_tgpp, trip_raw)

393 + trip_or_zero = in_range.select(trip_clamped, c_zero_i32)

394 + c_zero_idx = arith.constant(0, index=True)

395 + c_one_idx = arith.constant(1, index=True)

396 + trip_idx = arith.index_cast(

397 + T.index,

398 + trip_or_zero.ir_value()

399 + if hasattr(trip_or_zero, "ir_value")

400 + else trip_or_zero,

401 + )

402 + # Initial iter_args list: FA-2 accumulator state.

403 + # Order: running_max, running_sum, acc_pv[0..PV_N_CHUNKS-1].

404 + # NOTE: scf.ForOp requires ir.Value (with .type). The DSL

405 + # constants ZERO_F / ONE_F are fx.Float32 (Numeric) wrappers,

406 + # not raw ir.Value, so we unwrap via .ir_value() before

407 + # passing. _ival() is defined once near the top of the kernel body.

408 + _init_iter = [

409 + _ival(running_max),

410 + _ival(running_sum),

411 + *[_ival(p) for p in acc_pv],

412 + ]

413 + _for_op = _scf.ForOp(

414 + c_zero_idx,

415 + trip_idx,

416 + c_one_idx,

417 + _init_iter,

418 + )

419 + _for_ip = ir.InsertionPoint(_for_op.body)

420 + _for_ip.__enter__()

421 + try:

422 + tg_idx = _for_op.induction_variable

423 + tg_i32 = fx.Int32(arith.index_cast(T.i32, tg_idx))

424 + partition_start = partition_base + tg_i32 * c_kcb

425 + bt_seq_base = seq * c_bt + (partition_start // fx.Int32(_BS))

426 + running_max = _for_op.inner_iter_args[0]

427 + running_sum = _for_op.inner_iter_args[1]

428 + acc_pv = list(_for_op.inner_iter_args[2:])

429 + for n_tile in range_constexpr(16):

430 + block_in_part = n_tile // _TILES_PER_BLOCK

431 + tile_in_block = n_tile % _TILES_PER_BLOCK

432 + tile_start_tok = partition_start + fx.Int32(n_tile * TILE_SIZE)

433 + tile_in_seq = tile_start_tok < seq_len

434 + bt_off = bt_seq_base + fx.Int32(block_in_part)

435 + bt_off_safe = tile_in_seq.select(bt_off, seq * c_bt)

436 + phys_block = buffer_ops.buffer_load(

437 + bt_rsrc,

438 + bt_off_safe,

439 + vec_width=1,

440 + dtype=T.i32,

441 + )

442 + block_base = phys_block * c_block

443 + data_region = block_base

444 + meta_region = block_base + c_meta_off

445 +

446 + # ``slot`` is the absolute slot index within the cache block

447 + # (range 0.._BS-1). For BS=16 it equals tok_in_tile; for BS=32

448 + # it equals tile_in_block * 16 + tok_in_tile. Used both for

449 + # the K/V data region and SoA metadata region addressing.

450 + slot = fx.Int32(tile_in_block * TILE_SIZE) + tok_in_tile

451 + data_bases_byte = (

452 + data_region

453 + + slot * (c_kv_heads * c_data_per_slot)

454 + + kv_h * c_data_per_slot

455 + )

456 + k_byte = data_bases_byte + chunk_in_tok * fx.Int32(16)

457 + k_packed = buffer_ops.buffer_load(

458 + kv_rsrc,

459 + k_byte // fx.Int32(4),

460 + vec_width=4,

461 + dtype=T.i32,

462 + )

463 +

464 + # ---- HOISTED: issue V data + meta HBM loads early ---------

465 + # These are async; their s_waitcnt is pushed by the compiler

466 + # past the K dequant + QK MFMA + softmax block, hiding most

467 + # of the V HBM latency behind compute.

468 + v_byte = data_bases_byte + c_keydata + chunk_in_tok * fx.Int32(16)

469 + v_packed = buffer_ops.buffer_load(

470 + kv_rsrc,

471 + v_byte // fx.Int32(4),

472 + vec_width=4,

473 + dtype=T.i32,

474 + )

475 + vscale_u16 = (

476 + meta_region // fx.Int32(2)

477 + + kv_h * c_meta_u16_per_kvh

478 + + c_vscale_off_u16

479 + + slot

480 + )

481 + vzero_u16 = (

482 + meta_region // fx.Int32(2)

483 + + kv_h * c_meta_u16_per_kvh

484 + + c_vzero_off_u16

485 + + slot

486 + )

487 + vscale_raw = buffer_ops.buffer_load(

488 + kv_rsrc,

489 + vscale_u16,

490 + vec_width=1,

491 + dtype=T.i16,

492 + )

493 + vzero_raw = buffer_ops.buffer_load(

494 + kv_rsrc,

495 + vzero_u16,

496 + vec_width=1,

497 + dtype=T.i16,

498 + )

499 +

500 + knorm_u16 = (

501 + meta_region // fx.Int32(2)

502 + + kv_h * c_meta_u16_per_kvh

503 + + c_knorm_off_u16

504 + + slot

505 + )

506 + knorm_raw = buffer_ops.buffer_load(

507 + kv_rsrc,

508 + knorm_u16,

509 + vec_width=1,

510 + dtype=T.i16,

511 + )

512 + knorm_f16 = arith.bitcast(T.f16, knorm_raw)

513 + knorm_f32 = arith.extf(T.f32, knorm_f16)

514 +

515 + # K dequant → LDS [token, head_dim] (natural)

516 + # Lane writes 32 bf16 (= 4×8) for token=tok_in_tile,

517 + # head_dims chunk_in_tok*32..+31 (4 sub-chunks of 8).

518 + tok_kreg = tok_in_tile * fx.Int32(HEAD_SIZE * 2 // 8)

519 + chunk_kreg = tok_kreg + chunk_in_tok * fx.Int32(8)

520 + for w in range_constexpr(4):

521 + word_i32 = vector.extract(k_packed, static_position=[w])

522 + bf16_elems = []

523 + for n in range_constexpr(8):

524 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

525 + nibble_idx = arith.index_cast(T.index, nibble)

526 + cent_f32 = cent_lds.load([nibble_idx])

527 + elem_bf16 = arith.trunc_f(T.bf16, cent_f32 * knorm_f32)

528 + bf16_elems.append(elem_bf16)

529 + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems)

530 + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16)

531 + vector.store(

532 + v_i64,

533 + kv_lds_i64,

534 + [arith.index_cast(T.index, chunk_kreg + fx.Int32(w * 2))],

535 + )

536 + gpu.barrier()

537 +

538 + # QK MFMA (CDNA4 wide-K): A=K[token, head_dim], B=Q (= Q^T).

539 + # K read: lane t = K[token=mfma_row, head_dim=chunk*32+col_grp*8..+7]

540 + # Same i64×2 indexing as Q.

541 + qk_acc = zero_v4

542 + for chk in range_constexpr(QK_K_CHUNKS):

543 + k_idx_i64 = (

544 + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8)

545 + + fx.Int32(chk * 8)

546 + + mfma_col_grp * fx.Int32(2)

547 + )

548 + kv_load = vector.load_op(

549 + T.vec(2, T.i64),

550 + kv_lds_i64,

551 + [arith.index_cast(T.index, k_idx_i64)],

552 + )

553 + k_op = vector.bitcast(T.vec(8, T.bf16), kv_load)

554 + qk_acc = rocdl.mfma_f32_16x16x32_bf16(

555 + T.f32x4, [k_op, q_chunks[chk], qk_acc, 0, 0, 0]

556 + )

557 +

558 + # qk_acc layout: lane t holds C[token=(t/16)*4..+3, query=t%16]

559 + # 4 fp32/lane = 4 different tokens at SAME query_row.

560 +

561 + # Scale + mask out-of-context tokens.

562 + qk_acc = _vsplat_mul(qk_acc, QK_SCALE)

563 + for elem in range_constexpr(4):

564 + kv_tok = (

565 + partition_start

566 + + fx.Int32(n_tile * TILE_SIZE)

567 + + mfma_col_grp * fx.Int32(4)

568 + + fx.Int32(elem)

569 + )

570 + in_b = kv_tok < seq_len

571 + v = vector.extract(qk_acc, static_position=[elem])

572 + qk_acc = vector.insert(

573 + in_b.select(v, NEG_INF),

574 + qk_acc,

575 + static_position=[elem],

576 + dynamic_position=[],

577 + )

578 +

579 + # FA2 online softmax: per-query-row reduce.

580 + # Per-row max: max over 4 fp32 in lane, then xor-shuffle 16, 32

581 + # (across the 4 col_grps that share same mfma_row).

582 + local_max = vector.reduction(T.f32, "maxnumf", qk_acc)

583 + r1 = local_max.shuffle_xor(fx.Int32(16), c_w)

584 + local_max = local_max.maximumf(r1)

585 + r2 = local_max.shuffle_xor(fx.Int32(32), c_w)

586 + tile_max = local_max.maximumf(r2)

587 +

588 + new_max = running_max.maximumf(tile_max)

589 + max_diff = running_max - new_max

590 + safe_diff = (running_max > NEG_INF).select(max_diff, ZERO_F)

591 + scale = (safe_diff * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast)

592 + running_sum = running_sum * scale

593 + for h in range_constexpr(PV_N_CHUNKS):

594 + acc_pv[h] = _vsplat_mul(acc_pv[h], scale)

595 + running_max = new_max

596 +

597 + # Compute probs: p = exp((qk - new_max) * LOG2E)

598 + tile_sum = ZERO_F

599 + for elem in range_constexpr(4):

600 + s = vector.extract(qk_acc, static_position=[elem])

601 + d = s - new_max

602 + d = (new_max > NEG_INF).select(d, NEG_INF)

603 + p = (d * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast)

604 + tile_sum = tile_sum + p

605 + qk_acc = vector.insert(

606 + p, qk_acc, static_position=[elem], dynamic_position=[]

607 + )

608 +

609 + ts1 = tile_sum.shuffle_xor(fx.Int32(16), c_w)

610 + tile_sum = tile_sum + ts1

611 + ts2 = tile_sum.shuffle_xor(fx.Int32(32), c_w)

612 + tile_sum = tile_sum + ts2

613 + running_sum = running_sum + tile_sum

614 +

615 + # ---- V dequant → LDS (V/scale/zero raw loads were hoisted) -

616 + # v_packed, vscale_raw, vzero_raw are already in flight from

617 + # the top of this iteration; only the f32 conversion stays

618 + # here so the actual consume site is still close to the use.

619 + vscale_f32 = arith.extf(T.f32, arith.bitcast(T.f16, vscale_raw))

620 + vzero_f32 = arith.extf(T.f32, arith.bitcast(T.f16, vzero_raw))

621 +

622 + if USE_HW_TR:

623 + # V dequant → LDS [token][head_dim] (ROW-MAJOR, no transpose).

624 + # Each lane writes 32 contiguous bf16 (one token, head_dims

625 + # chunk_in_tok*32..+31) as 4× ds_write_b128 = 4× vec(2,i64).

626 + # Replaces 32× ds_write_b16 of the legacy transposed path.

627 + v_lds_elem_base = tok_in_tile * fx.Int32(

628 + HEAD_SIZE

629 + ) + chunk_in_tok * fx.Int32(32)

630 + for w in range_constexpr(4):

631 + word_i32 = vector.extract(v_packed, static_position=[w])

632 + bf16_elems = []

633 + for n in range_constexpr(8):

634 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

635 + nibble_f32 = arith.sitofp(T.f32, nibble)

636 + elem_f32 = nibble_f32 * vscale_f32 + vzero_f32

637 + elem_bf16 = arith.trunc_f(T.bf16, elem_f32)

638 + bf16_elems.append(elem_bf16)

639 + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems)

640 + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16)

641 + v_lds_i64_idx = (v_lds_elem_base + fx.Int32(w * 8)) // fx.Int32(

642 + 4

643 + )

644 + vector.store(

645 + v_i64,

646 + kv_lds_i64,

647 + [arith.index_cast(T.index, v_lds_i64_idx)],

648 + )

649 + # ── HW V transpose: cross-lane LDS sync ─────────────────────

650 + # ds_read_tr16_b64 below introduces a cross-lane LDS read:

651 + # consumer lane t reads bytes written by source lanes

652 + # 0,4,8,12 (etc.). The compiler's automatic waitcnt insertion

653 + # is per-lane and has no awareness of the HW transpose's

654 + # cross-lane forwarding, so it can sink the next iteration's

655 + # ds_read ahead of these ds_write_b128 ops in the schedule.

656 + #

657 + # Two-part fence:

658 + # 1. sched_barrier(0): MachineScheduler reorder barrier

659 + # (mask=0 → no instruction class may cross). Pure hint,

660 + # generates no runtime instruction. Use the simple

661 + # sched_barrier (NOT sched_group_barrier, which is an

662 + # IGLP marker that requires partner barriers and

663 + # corrupts the schedule when used in isolation).

664 + # 2. s_waitcnt lgkmcnt=0: runtime drain of the LDS write

665 + # queue. vmcnt/expcnt left at no-wait so we don't

666 + # stall on speculatively-issued HBM loads (which can

667 + # include OOB-but-masked buffer_load reads from the

668 + # next K-tile's hoisted V prefetch).

669 + #

670 + # The SW path below does NOT need this because each lane

671 + # only reads cells it itself wrote (no cross-lane dep), and

672 + # the per-lane same-address waitcnt the compiler emits is

673 + # correct.

674 + #

675 + # Without this fence the reordered ds_read produces

676 + # wrong results on high-partition-count configs (large

677 + # padded num_partitions); low-partition configs are

678 + # unaffected.

679 + rocdl.sched_barrier(0)

680 + # encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=0)

681 + # = 0xF | (7<<4) | (0<<8) | (3<<14) = 0xC07F

682 + rocdl.s_waitcnt(0xC07F)

683 + else:

684 + # Legacy transposed V_LDS path: 32 ds_write_b16 per lane.

685 + for w in range_constexpr(4):

686 + word_i32 = vector.extract(v_packed, static_position=[w])

687 + for n in range_constexpr(8):

688 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

689 + nibble_f32 = arith.sitofp(T.f32, nibble)

690 + elem_f32 = nibble_f32 * vscale_f32 + vzero_f32

691 + elem_bf16 = arith.trunc_f(T.bf16, elem_f32)

692 + elem_i16 = arith.bitcast(T.i16, elem_bf16)

693 + head_dim = chunk_in_tok * fx.Int32(32) + fx.Int32(w * 8 + n)

694 + v_idx_i16 = head_dim * fx.Int32(TILE_SIZE) + tok_in_tile

695 + v_vec = vector.from_elements(T.vec(1, T.i16), [elem_i16])

696 + vector.store(

697 + v_vec,

698 + kv_lds_i16,

699 + [arith.index_cast(T.index, v_idx_i16)],

700 + )

701 + gpu.barrier()

702 +

703 + # ---- PV MFMA: A=V[head_dim, token], B=P (=qk_acc bf16) -----

704 + # P operand B layout matches qk_acc: lane t holds 4 bf16 at

705 + # K=token=(t/16)*4..+3, N=query=t%16. Just trunc_f to bf16.

706 + p_bf16 = arith.trunc_f(T.vec(4, T.bf16), qk_acc)

707 + p_op = vector.bitcast(T.vec(4, T.i16), p_bf16)

708 +

709 + if USE_HW_TR:

710 + # HW-transpose PV path: V_lds is row-major V[token][head_dim].

711 + # MFMA A operand layout: lane t holds 4 bf16 at

712 + # M=head_dim=mfma_row + h*16, K=token=(t/16)*4..+3

713 + # ds_read_tr16_b64 (4-element 16-bit transpose, per-16-lane block):

714 + # result[lane=t, elem=e] = Input[source_lane=e*4 + (t%16)//4,

715 + # col=t%4]

716 + # Per-lane address: token_idx = lane // 4 (covers 0..15 across

717 + # all four 16-lane MFMA blocks), hd_sub = (lane % 4)*4 selects

718 + # the 4-element column window inside the h*16 chunk.

719 + # Total LDS byte offset = kv_off + token_idx*HEAD_SIZE*2

720 + # + (h*16 + hd_sub)*2

721 + token_idx = lane >> fx.Int32(2)

722 + hd_sub = (lane & fx.Int32(3)) * fx.Int32(4)

723 + v_lane_byte = (

724 + fx.Int32(kv_off)

725 + + token_idx * fx.Int32(HEAD_SIZE * 2)

726 + + hd_sub * fx.Int32(2)

727 + )

728 + for h in range_constexpr(PV_N_CHUNKS):

729 + v_byte_off = v_lane_byte + fx.Int32(h * 32)

730 + v_byte_i64 = fx.Int64(v_byte_off)

731 + v_ptr = buffer_ops.create_llvm_ptr(

732 + v_byte_i64,

733 + address_space=3,

734 + )

735 + v_op_raw = rocdl.ds_read_tr16_b64(

736 + T.vec(4, T.i16),

737 + v_ptr,

738 + ).result

739 + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k(

740 + T.f32x4, [v_op_raw, p_op, acc_pv[h], 0, 0, 0]

741 + )

742 + else:

743 + # Legacy transposed-LDS PV: 1 ds_read_b64 per chunk.

744 + # Byte addr = (mfma_row + h*16) * TILE_SIZE * 2 + (col_grp*4) * 2

745 + # = (mfma_row + h*16) * 32 + col_grp*8

746 + # i64 idx = (mfma_row + h*16) * 4 + col_grp

747 + for h in range_constexpr(PV_N_CHUNKS):

748 + v_idx_i64 = (mfma_row + fx.Int32(h * 16)) * fx.Int32(

749 + 4

750 + ) + mfma_col_grp

751 + kv_load = vector.load_op(

752 + T.vec(1, T.i64),

753 + kv_lds_i64,

754 + [arith.index_cast(T.index, v_idx_i64)],

755 + )

756 + v_op = vector.bitcast(T.vec(4, T.i16), kv_load)

757 + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k(

758 + T.f32x4, [v_op, p_op, acc_pv[h], 0, 0, 0]

759 + )

760 +

761 + _scf.YieldOp(

762 + [

763 + _ival(running_max),

764 + _ival(running_sum),

765 + *[_ival(p) for p in acc_pv],

766 + ]

767 + )

768 + finally:

769 + _for_ip.__exit__(None, None, None)

770 + # Pull final accumulator values out of the for_op results.

771 + running_max = _for_op.results[0]

772 + running_sum = _for_op.results[1]

773 + acc_pv = list(_for_op.results[2:])

774 +

775 + # ===== STEP G: Output ===========================================

776 + # acc_pv[h] layout: lane t holds 4 fp32 at M=head_dim=(t/16)*4..+3 + h*16,

777 + # N=query_row=t%16.

778 + # 4 fp32/lane = 4 contiguous head_dim values for ONE query_row.

779 + safe_sum = (running_sum > ZERO_F).select(running_sum, ONE_F)

780 + rcp = ONE_F / safe_sum

781 +

782 + c_os = fx.Int32(_stride_out_seq)

783 + c_oh = fx.Int32(_stride_out_head)

784 + c_op_ = fx.Int32(_stride_out_part)

785 + out_base = seq * c_os + kv_h * c_oh + part * c_op_

786 +

787 + # Gate all global stores by mfma_row < QG. For QG=16 the predicate is

788 + # tautological (mfma_row ∈ [0,15]); for QG=8 lanes 8..15 (whose MFMA

789 + # outputs are computational waste) skip the stores entirely.

790 + valid_row_pred = arith.cmpi(

791 + arith.CmpIPredicate.ult,

792 + mfma_row.ir_value() if hasattr(mfma_row, "ir_value") else mfma_row,

793 + arith.constant(QG, type=T.i32),

794 + )

795 + _if = _scf.IfOp(valid_row_pred)

796 + with ir.InsertionPoint(_if.then_block):

797 + for h in range_constexpr(PV_N_CHUNKS):

798 + pv_norm = _vsplat_mul(acc_pv[h], rcp)

799 + pv_bf16 = arith.trunc_f(T.vec(4, T.bf16), pv_norm)

800 + pv_i32x2 = vector.bitcast(T.vec(2, T.i32), pv_bf16)

801 + head_dim_start = fx.Int32(h * 16) + mfma_col_grp * fx.Int32(4)

802 + out_off_elem = (

803 + out_base + mfma_row * fx.Int32(HEAD_SIZE) + head_dim_start

804 + )

805 + buffer_ops.buffer_store(

806 + pv_i32x2,

807 + out_rsrc,

808 + out_off_elem * fx.Int32(2),

809 + offset_is_bytes=True,

810 + )

811 +

812 + c_npq = fx.Int32(num_partitions * QG)

813 + ml_off = (

814 + seq * fx.Int32(_stride_ml_seq) + kv_h * c_npq + part * c_qg + mfma_row

815 + )

816 + es_off = (

817 + seq * fx.Int32(_stride_es_seq) + kv_h * c_npq + part * c_qg + mfma_row

818 + )

819 + buffer_ops.buffer_store(running_max, ml_rsrc, ml_off)

820 + buffer_ops.buffer_store(running_sum, es_rsrc, es_off)

821 + _scf.YieldOp([])

822 +

823 + return tq_decode_kernel

vllm/v1/attention/ops/flydsl_kernels/tq_decode_gqa6.py new

+875

@@ -0,0 +1,875 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# TurboQuant decode GQA-6 sibling (FlyDSL) — MI355X / gfx950 / CDNA4

4 +#

5 +# Decode-only paged attention over a TurboQuant SoA KV cache.

6 +# Targets MiniMax-M2.5 class models: HEAD_SIZE=128, GQA group=6, MSE_BITS=4 K,

7 +# VQB=4 V, N_CENTROIDS=16, BLOCK_SIZE=16, partitioned mode.

8 +#

9 +# === Sibling-file rationale ===

10 +# Structurally identical to tq_decode.py except:

11 +# * QUERY_GROUP_SIZE default = 6 (was 8)

12 +# * Build-time assert allows query_group_size = 6 only (Qwen sizes go to

13 +# the canonical tq_decode kernel — NOT this file)

14 +# * QG_LOAD_ITERS = (QG + 3) // 4 (round-up; 2 iters for QG=6 loads

15 +# 8 LDS rows of which only 6 carry real Q; rows 6..7 are zero-padded

16 +# by buffer_load OOB protection — gated out by the same mfma_row < QG

17 +# predicate the canonical kernel already uses for QG=8)

18 +#

19 +# All cross-lane operations (xor_shuffle softmax reductions, ds_read_tr16_b64

20 +# HW V transpose) are mfma_row-isolated and thus already GQA-agnostic. The

21 +# kernel is functionally a direct copy of tq_decode.py with the QG knob

22 +# turned to 6 — kept as a sibling per FlyDSL convention to preserve the

23 +# canonical kernel's invariants for its production callers (Qwen2.5-72B,

24 +# Qwen3-32B).

25 +#

26 +# === Design ===

27 +# * 1 wave (64 lanes) per CTA — simpler than 4-warp; trades parallelism for

28 +# correctness/debug-ability. A later revision will fan out to 4 warps.

29 +# * Flash-Attention-2 online softmax: rescale running PV/sum per K-tile,

30 +# no QK score storage across tiles, no P-LDS round-trip (qk_acc layout

31 +# matches PV B-operand directly).

32 +# * Centroids loaded once into LDS at kernel start (LDS-resident LUT).

33 +# * K dequant → LDS [token, head_dim] (natural layout)

34 +# * V dequant → LDS [head_dim, token] (TRANSPOSED — required so PV MFMA's

35 +# A=V operand can be loaded as 4-bf16 contiguous via 1 ds_read_b64).

36 +#

37 +# MFMA layouts used:

38 +# QK: mfma(A=K, B=Q, C=qk_acc) → C[m=token, n=query], 4 fp32/lane

39 +# = 4 tokens for SAME query → softmax via

40 +# xor_shuffle(16),(32) reduces per-row

41 +# PV: mfma(A=V_T, B=P, C=acc_pv) → C[m=head_dim, n=query], 4 fp32/lane

42 +# = 4 contiguous head_dim cols for SAME

43 +# query → 1 buffer_store of 8 bytes

44 +# per chunk per lane

45 +

46 +from __future__ import annotations

47 +

48 +import flydsl.compiler as flyc

49 +import flydsl.expr as fx

50 +from flydsl._mlir import ir

51 +from flydsl._mlir.dialects import scf as _scf

52 +from flydsl.expr import (

53 + arith,

54 + buffer_ops,

55 + const_expr, # noqa: F401 kept for kernel-DSL API parity

56 + gpu,

57 + range_constexpr,

58 + rocdl,

59 + vector,

60 +)

61 +from flydsl.expr.typing import Int32, T # noqa: F401 Int32 via fx.Int32

62 +from flydsl.runtime.device import get_rocm_arch as get_hip_arch

63 +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr

64 +

65 +# === Constants (MiniMax-M2.5-class TQ decode profile) =======================

66 +HEAD_SIZE = 128

67 +# default; overridable via build_tq_decode_gqa6_module(kv_block_size=...)

68 +KV_BLOCK_SIZE = 16

69 +# MFMA tile = 16 tokens (do not change without re-deriving MFMA shapes)

70 +TILE_SIZE = 16

71 +N_CENTROIDS = 16

72 +QUERY_GROUP_SIZE = 6 # GQA-6 default for MiniMax-M2.5

73 +WARP_SIZE = 64

74 +NUM_WARPS = 1

75 +BLOCK_THREADS = NUM_WARPS * WARP_SIZE

76 +KV_COMPUTE_BLOCK = 256 # 16 K-tiles × 16 tokens

77 +

78 +# TQ SoA layout

79 +NUM_SOA_FIELDS = 3

80 +SOA_K_NORM = 0

81 +SOA_V_SCALE = 1

82 +SOA_V_ZERO = 2

83 +KEY_DATA_BYTES = HEAD_SIZE // 2

84 +VAL_DATA_BYTES = HEAD_SIZE // 2

85 +DATA_BYTES_PER_SLOT = KEY_DATA_BYTES + VAL_DATA_BYTES

86 +

87 +# MFMA

88 +MFMA_M = MFMA_N = 16

89 +MFMA_K_BF16_QK = 32 # CDNA4 wide-K (mfma_f32_16x16x32_bf16)

90 +MFMA_K_BF16_PV = 16 # PV's K = token tile = 16

91 +QK_K_CHUNKS = HEAD_SIZE // MFMA_K_BF16_QK # 4

92 +PV_N_CHUNKS = HEAD_SIZE // MFMA_N # 8

93 +

94 +# LDS regions

95 +# Q region uses the full 16-row footprint (4096 B) rather than a tight

96 +# QG=6 sizing, keeping the LDS layout identical to the canonical kernel.

97 +CENTROID_LDS_BYTES = N_CENTROIDS * 4 # 64

98 +Q_LDS_BYTES = 16 * HEAD_SIZE * 2 # 4096

99 +KV_TILE_LDS_BYTES = TILE_SIZE * HEAD_SIZE * 2 # 4096

100 +

101 +LOG2E = 1.4426950408889634

102 +NEG_INF_VAL = float("-inf")

103 +

104 +

105 +def _vsplat_mul(vec, scalar):

106 + s = scalar.ir_value() if hasattr(scalar, "ir_value") else scalar

107 + return vec * vector.broadcast(T.f32x4, s)

108 +

109 +

110 +allocator = None

111 +

112 +

113 +def build_tq_decode_gqa6_module(

114 + num_seqs: int,

115 + num_kv_heads: int,

116 + num_partitions: int,

117 + max_blocks_per_seq: int = 512,

118 + softmax_scale: float | None = None,

119 + query_group_size: int = QUERY_GROUP_SIZE,

120 + kv_block_size: int = KV_BLOCK_SIZE,

121 + use_hw_v_transpose: bool = False,

122 + tile_groups_per_partition: int = 1,

123 +):

124 + """Build a TQ decode GQA-6 kernel module (MiniMax-M2.5).

125 +

126 + ``query_group_size`` (= num_query_heads // num_kv_heads) is fixed at 6

127 + for this sibling. The MFMA's 16-row capacity is 6/16 = 37.5% used;

128 + lanes 6..15 compute garbage and are gated out of all global-memory

129 + writes via the same ``mfma_row < QG`` predicate the canonical kernel

130 + uses for QG=8. Q load iterates ``(QG + 3) // 4`` = 2 chunks (loads

131 + 8 LDS rows of which rows 6..7 are zero-padded by the buffer-resource

132 + OOB protection — they're never consumed because the gating lanes

133 + don't write).

134 +

135 + For GQA-8 or GQA-16 (Qwen models), use the canonical

136 + ``build_tq_decode_module`` from ``tq_decode.py`` instead.

137 +

138 + ``kv_block_size`` is the number of tokens per vLLM cache block. Must

139 + be a multiple of TILE_SIZE (=16). Supported values: 16 and 32.

140 +

141 + ``use_hw_v_transpose`` (CDNA4 / gfx950 only) — see canonical kernel

142 + docstring for details. Same fence pattern applies here.

143 + """

144 + assert query_group_size == 6, (

145 + f"build_tq_decode_gqa6_module is GQA-6 only; got "

146 + f"query_group_size={query_group_size}. Use build_tq_decode_module "

147 + f"from tq_decode.py for GQA-8 or GQA-16."

148 + )

149 + assert kv_block_size in (16, 32), (

150 + f"kv_block_size must be 16 or 32; got {kv_block_size}"

151 + )

152 + assert kv_block_size % TILE_SIZE == 0

153 + assert int(tile_groups_per_partition) >= 1, (

154 + f"tile_groups_per_partition must be >= 1; got {tile_groups_per_partition}"

155 + )

156 + USE_HW_TR = bool(use_hw_v_transpose)

157 + TGPP = int(tile_groups_per_partition)

158 + PARTITION_EXTENT_TOKENS = TGPP * KV_COMPUTE_BLOCK

159 +

160 + QG = int(query_group_size)

161 + # Round-up: 2 iters for QG=6 → loads 8 LDS rows (rows 6..7 are

162 + # zero-padded by buffer_load OOB protection; never consumed because

163 + # mfma_row >= QG lanes are gated out of all global stores below).

164 + QG_LOAD_ITERS = (QG + 3) // 4

165 + OOB_OFFSET = 0x7FFFFFF0 # noqa: F841 ~2GB byte offset; > any plausible buffer

166 +

167 + _BS = int(kv_block_size)

168 + _TILES_PER_BLOCK = _BS // TILE_SIZE # 1 for BS=16, 2 for BS=32

169 +

170 + global allocator

171 + arch = get_hip_arch()

172 +

173 + if softmax_scale is None:

174 + softmax_scale = 1.0 / (HEAD_SIZE**0.5)

175 + _qk_scale = float(softmax_scale)

176 +

177 + # --- Strides ---

178 + _Hq = num_kv_heads * QG

179 + _stride_q_seq = _Hq * HEAD_SIZE

180 + _stride_q_head = HEAD_SIZE

181 + _stride_bt_seq = max_blocks_per_seq

182 +

183 + _data_region_bytes = _BS * num_kv_heads * DATA_BYTES_PER_SLOT

184 + _meta_region_bytes = num_kv_heads * NUM_SOA_FIELDS * _BS * 2

185 + _stride_cache_block = _data_region_bytes + _meta_region_bytes

186 + _meta_region_offset = _data_region_bytes

187 +

188 + _stride_out_part = QG * HEAD_SIZE

189 + _stride_out_head = num_partitions * QG * HEAD_SIZE

190 + _stride_out_seq = num_kv_heads * num_partitions * QG * HEAD_SIZE

191 + _stride_es_seq = num_kv_heads * num_partitions * QG

192 + _stride_ml_seq = _stride_es_seq

193 +

194 + # --- LDS layout ---

195 + allocator = SmemAllocator(None, arch=arch, global_sym_name="tq_gqa6_smem")

196 + centroid_off = 0

197 + allocator.ptr = CENTROID_LDS_BYTES

198 + q_off = allocator.ptr

199 + allocator.ptr += Q_LDS_BYTES

200 + kv_off = allocator.ptr

201 + allocator.ptr += KV_TILE_LDS_BYTES

202 +

203 + @flyc.kernel

204 + def tq_decode_gqa6_kernel(

205 + out_ptr: fx.Tensor,

206 + exp_sums_ptr: fx.Tensor,

207 + max_logits_ptr: fx.Tensor,

208 + query_ptr: fx.Tensor,

209 + kv_cache_ptr: fx.Tensor,

210 + centroids_ptr: fx.Tensor,

211 + block_tables_ptr: fx.Tensor,

212 + seq_lens_ptr: fx.Tensor,

213 + ):

214 + # ---- IDs ---------------------------------------------------------

215 + tid = gpu.thread_idx.x

216 + seq = gpu.block_idx.x

217 + kv_h = gpu.block_idx.y

218 + part = gpu.block_idx.z

219 + lane = tid # 0..63

220 + mfma_row = lane & fx.Int32(15) # = query_row (B's N) /

221 + # token (A's M for QK)

222 + mfma_col_grp = lane >> fx.Int32(4) # 0..3, K-group dim

223 +

224 + # ---- Buffer resources -------------------------------------------

225 + q_rsrc = buffer_ops.create_buffer_resource(query_ptr, max_size=True)

226 + kv_rsrc = buffer_ops.create_buffer_resource(kv_cache_ptr, max_size=True)

227 + bt_rsrc = buffer_ops.create_buffer_resource(block_tables_ptr, max_size=True)

228 + sl_rsrc = buffer_ops.create_buffer_resource(seq_lens_ptr, max_size=True)

229 + cent_rsrc = buffer_ops.create_buffer_resource(centroids_ptr, max_size=True)

230 + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True)

231 + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True)

232 + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True)

233 +

234 + # ---- LDS pointers -----------------------------------------------

235 + base = allocator.get_base()

236 + cent_lds = SmemPtr(base, centroid_off, T.f32, shape=(N_CENTROIDS,))

237 + q_lds_i32 = SmemPtr(base, q_off, T.i32, shape=(Q_LDS_BYTES // 4,)).get()

238 + q_lds_i64 = SmemPtr(base, q_off, T.i64, shape=(Q_LDS_BYTES // 8,)).get()

239 + kv_lds_i32 = SmemPtr(base, kv_off, T.i32, shape=(KV_TILE_LDS_BYTES // 4,)).get() # noqa: F841

240 + kv_lds_i64 = SmemPtr(base, kv_off, T.i64, shape=(KV_TILE_LDS_BYTES // 8,)).get()

241 + kv_lds_i16 = SmemPtr(base, kv_off, T.i16, shape=(KV_TILE_LDS_BYTES // 2,)).get()

242 +

243 + # ---- Constants ---------------------------------------------------

244 + c_sq = fx.Int32(_stride_q_seq)

245 + c_qh = fx.Int32(_stride_q_head)

246 + c_qg = fx.Int32(QG)

247 + c_bt = fx.Int32(_stride_bt_seq)

248 + c_block = fx.Int32(_stride_cache_block)

249 + c_meta_off = fx.Int32(_meta_region_offset)

250 + c_data_per_slot = fx.Int32(DATA_BYTES_PER_SLOT)

251 + c_kv_heads = fx.Int32(num_kv_heads)

252 + c_keydata = fx.Int32(KEY_DATA_BYTES)

253 + c_w = fx.Int32(WARP_SIZE)

254 +

255 + NEG_INF = arith.constant(NEG_INF_VAL, type=T.f32)

256 + ZERO_F = fx.Float32(0.0)

257 + ONE_F = fx.Float32(1.0)

258 + LOG2E_C = arith.constant(LOG2E, type=T.f32)

259 + QK_SCALE = arith.constant(_qk_scale, type=T.f32)

260 +

261 + # ===== STEP A: Load centroids → LDS (cooperative, race-safe) =====

262 + c_idx_safe = lane & fx.Int32(N_CENTROIDS - 1)

263 + c_val = buffer_ops.buffer_load(cent_rsrc, c_idx_safe, vec_width=1, dtype=T.f32)

264 + cent_lds.store(c_val, [arith.index_cast(T.index, c_idx_safe)])

265 + gpu.barrier()

266 +

267 + # ===== STEP B: Load Q → row-major LDS [query_row, head_dim] ======

268 + # Q layout in HBM: [N, Hq, D] bf16, slice = QG rows × D=128 cols × 2B.

269 + # QG=6 → 1536 B real, 96 chunks; QG_LOAD_ITERS=2 issues 128 chunks

270 + # (8 rows × 16 col_b) so lanes whose row ∈ {6, 7} address heads that

271 + # do NOT exist in the Q tensor (Hq = num_kv_heads * 6).

272 + #

273 + # ── Per-lane Q-load OOB redirect ─────────────────────────────────

274 + # Hq = num_kv_heads * QG, so head index = kv_h * QG + row. For QG=6

275 + # the upper rows alias to (kv_h+1)'s real heads (still in-bounds

276 + # for most kv_h), but for the *last* (seq, kv_h) the read tips

277 + # past the Q tensor entirely. With ``max_size=True`` the buffer

278 + # descriptor advertises 4 GB so the HW does no OOB clamping; the

279 + # padding read then touches unmapped pages and the kernel dies

280 + # with "Memory access fault by GPU node" the moment the

281 + # surrounding allocator state leaves the next page unmapped

282 + # (allocator-pattern-dependent → nondeterministic fault).

283 + #

284 + # Fix: redirect OOB lanes (``row >= QG``) to load Q[0,0,col_elem]

285 + # — a byte offset that is always inside the actual Q tensor.

286 + # The redundant load is wasted but harmless; the resulting LDS

287 + # write to row ∈ {6, 7} is overwritten/ignored downstream because

288 + # every global store is gated by the output-time

289 + # ``mfma_row < QG`` predicate (the same gate canonical QG=8 uses

290 + # for its garbage rows 8..15).

291 + #

292 + # NB: A ``mask=`` parameter on ``buffer_load`` substitutes the

293 + # masked-out offset with 0x7FFFFFFF *bytes*; with ``max_size=True``

294 + # this is still inside the 4 GB descriptor and so the HW happily

295 + # dereferences ``q_base + 2 GB`` → another unmapped page → fault.

296 + # The address-redirect below is the only safe option short of

297 + # threading the true Q size through as a runtime kernel arg.

298 + #

299 + # The canonical Qwen kernel does not need this redirect because

300 + # for QG ∈ {8, 16} the (kv_h, row) iteration covers exactly the

301 + # real heads; QG=6 is the only case where

302 + # ``QG_LOAD_ITERS * 4 != QG``.

303 + # STEP B: load pre-rotated q_rot → Q_LDS. QG=6 over-reads

304 + # (QG_LOAD_ITERS * 4 != QG), so rows >= QG redirect to a safe

305 + # in-bounds offset to avoid faulting past the real Q rows.

306 + for c in range_constexpr(QG_LOAD_ITERS):

307 + row_chunk = lane + fx.Int32(c * WARP_SIZE) # 0..127

308 + row = row_chunk >> fx.Int32(4) # 0..7

309 + col_b = row_chunk & fx.Int32(15) # 0..15

310 + col_elem = col_b * fx.Int32(8) # 0..120 in bf16

311 + q_off_real = seq * c_sq + (kv_h * c_qg + row) * c_qh + col_elem

312 + row_in_qg = row < fx.Int32(QG)

313 + q_off_elem = row_in_qg.select(q_off_real, col_elem)

314 + q_v = buffer_ops.buffer_load(

315 + q_rsrc,

316 + q_off_elem // fx.Int32(2),

317 + vec_width=4,

318 + dtype=T.i32,

319 + )

320 + q_lds_byte = row * fx.Int32(HEAD_SIZE * 2) + col_elem * fx.Int32(2)

321 + vector.store(

322 + q_v,

323 + q_lds_i32,

324 + [arith.index_cast(T.index, q_lds_byte // fx.Int32(4))],

325 + )

326 + gpu.barrier()

327 +

328 + # ===== STEP C: Pre-load Q operands for QK_K_CHUNKS = 4 K-chunks ===

329 + # Wide-K MFMA: K_per_chunk = 32, K_per_lane = 32/4 = 8 bf16 per chunk.

330 + # Lane t holds 8 bf16 at row=mfma_row, K-cols

331 + # chunk*32 + mfma_col_grp*8 + 0..7

332 + # Byte addr = mfma_row * HEAD_SIZE*2 + (chunk*32 + col_grp*8)*2

333 + # = mfma_row * 256 + chunk*64 + col_grp*16

334 + # i64 idx = mfma_row * 32 + chunk*8 + col_grp*2 (each i64 = 4 bf16)

335 + # Load 16 bytes (= 8 bf16 = 2 i64) per chunk via vec(2, i64).

336 + q_chunks = []

337 + for chk in range_constexpr(QK_K_CHUNKS):

338 + q_idx_i64 = (

339 + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8)

340 + + fx.Int32(chk * 8)

341 + + mfma_col_grp * fx.Int32(2)

342 + )

343 + qv = vector.load_op(

344 + T.vec(2, T.i64),

345 + q_lds_i64,

346 + [arith.index_cast(T.index, q_idx_i64)],

347 + )

348 + q_chunks.append(vector.bitcast(T.vec(8, T.bf16), qv))

349 +

350 + # ===== STEP D: Online softmax + PV state =========================

351 + running_max = NEG_INF

352 + running_sum = ZERO_F

353 + zero_v4 = arith.constant_vector(0.0, T.f32x4)

354 + acc_pv = [zero_v4 for _ in range(PV_N_CHUNKS)]

355 +

356 + # ===== STEP E: Sequence-len + partition base ====================

357 + seq_len = buffer_ops.buffer_load(sl_rsrc, seq, vec_width=1, dtype=T.i32)

358 + # ── Bounded num_partitions + internal looping ───────────────────

359 + # Mirrors the canonical tq_decode kernel: each CTA owns

360 + # ``PARTITION_EXTENT_TOKENS = TGPP * KV_COMPUTE_BLOCK`` contiguous

361 + # tokens of the seq, processed as TGPP groups of 16 K-tiles each.

362 + # FA-2 online-softmax state accumulates across all TGPP * 16 tiles.

363 + # With TGPP=1 this is bit-identical to the legacy 16-tile body.

364 + partition_base = part * fx.Int32(PARTITION_EXTENT_TOKENS)

365 +

366 + # Per-K-tile dequant lane assignment:

367 + # 16 tokens × 64 packed bytes per token = 1024 bytes

368 + # 64 lanes × 16 bytes/lane = 1024 ✓

369 + # Lane t: token = t / 4, sub-chunk = t % 4 (0..3) → 16 bytes contiguous

370 + tok_in_tile = lane >> fx.Int32(2)

371 + chunk_in_tok = lane & fx.Int32(3)

372 + # ``slot`` (= absolute slot index within the cache block, range

373 + # 0.._BS-1) is recomputed per tile inside the K-tile loop below

374 + # so it reflects ``tile_in_block * TILE_SIZE + tok_in_tile``.

375 +

376 + c_meta_u16_per_kvh = fx.Int32(NUM_SOA_FIELDS * _BS)

377 + c_knorm_off_u16 = fx.Int32(SOA_K_NORM * _BS)

378 + c_vscale_off_u16 = fx.Int32(SOA_V_SCALE * _BS)

379 + c_vzero_off_u16 = fx.Int32(SOA_V_ZERO * _BS)

380 +

381 + # ===== STEP F: K-tile loop =======================================

382 + # With kv_block_size > TILE_SIZE there are _TILES_PER_BLOCK tiles

383 + # per cache block. Tiles within the same block share a block-table

384 + # entry but address different slot ranges within the block.

385 + #

386 + # ── Per-tile block-table OOB redirect ────────────────────────────

387 + # The K-tile loop is unrolled 16 times so every CTA issues 16

388 + # block-table reads regardless of (partition, seq_len). When a

389 + # partition extends past ``seq_len`` (e.g. num_partitions=2 for a

390 + # 256-token sequence — partition 1 covers tokens 256..511, all

391 + # masked out) the per-tile bt offset can land beyond the bt

392 + # allocation. With ``max_size=True`` the descriptor is 4 GB so

393 + # the HW returns whatever pre-existing HBM bytes live at that

394 + # offset; the resulting garbage ``phys_block`` is multiplied by

395 + # ``stride_cache_block`` and the subsequent kv_cache read jumps

396 + # into unmapped pages, killing the whole launch with a

397 + # "Memory access fault by GPU node" exactly as the OOB Q load

398 + # used to (same allocator-pattern-dependent fault; B=4 makes

399 + # it deterministic by adding more colliding CTAs).

400 + #

401 + # Fix: when the tile starts at or past ``seq_len`` (so its

402 + # qk_acc is going to be killed by the per-token

403 + # ``kv_tok < seq_len`` mask anyway), redirect the bt read to

404 + # ``bt[seq, 0]`` — always in bounds. The redundant phys_block

405 + # decode + kv_cache read is wasted work, but correctness is

406 + # preserved without changing the kernel's iteration count.

407 + #

408 + # The canonical Qwen kernel reads the same tiles and is also

409 + # exposed to this issue in principle; in practice it surfaces on

410 + # GQA-6 first because the surrounding allocator state differs

411 + # (extra Q-tail allocation) and shifts the page maps so the bt

412 + # overflow lands on unmapped pages. Applying the same redirect to

413 + # the canonical kernel is a possible follow-up.

414 + # ===== Runtime-adaptive scf.ForOp w/ iter_args =========

415 + # Runtime-adaptive outer loop: trip count derives from

416 + # actual seq_len, NOT TGPP_max. At cudagraph-capture warmup

417 + # (seq_len=1) only 1 iteration runs; at 32K production

418 + # decode all TGPP iterations run. Kernel binary stays small

419 + # (single body, looped at runtime) — capture time matches

420 + # the legacy 16-tile-only kernel rather than scaling 4x.

421 + #

422 + # FA-2 state (running_max, running_sum, acc_pv[8]) threads

423 + # through scf iter_args; the body reads them at top, runs

424 + # the unchanged 16-tile inner body, and yields the new

425 + # state at bottom. After the loop, results are pulled out

426 + # of for_op.results back into the local Python names so the

427 + # downstream STEP G (output) is unmodified.

428 + c_kcb = fx.Int32(KV_COMPUTE_BLOCK)

429 + c_tgpp = fx.Int32(TGPP)

430 + c_zero_i32 = fx.Int32(0)

431 + c_one_i32 = fx.Int32(1)

432 + # remaining = seq_len - partition_base (signed, may be <=0)

433 + remaining = seq_len - partition_base

434 + in_range = remaining > c_zero_i32

435 + # trip_raw = ceil(remaining / KV_COMPUTE_BLOCK)

436 + # (when in_range is false, the divisor branch produces

437 + # garbage that the select below overrides with zero, so we

438 + # don't pre-clamp)

439 + trip_raw = (remaining + c_kcb - c_one_i32) // c_kcb

440 + trip_clamped = (trip_raw > c_tgpp).select(c_tgpp, trip_raw)

441 + trip_or_zero = in_range.select(trip_clamped, c_zero_i32)

442 + c_zero_idx = arith.constant(0, index=True)

443 + c_one_idx = arith.constant(1, index=True)

444 + trip_idx = arith.index_cast(

445 + T.index,

446 + trip_or_zero.ir_value()

447 + if hasattr(trip_or_zero, "ir_value")

448 + else trip_or_zero,

449 + )

450 +

451 + # Initial iter_args list: FA-2 accumulator state.

452 + # Order: running_max, running_sum, acc_pv[0..PV_N_CHUNKS-1].

453 + # NOTE: scf.ForOp requires ir.Value (with .type). The DSL

454 + # constants ZERO_F / ONE_F are fx.Float32 (Numeric) wrappers,

455 + # not raw ir.Value, so we unwrap via .ir_value() before

456 + # passing.

457 + def _ival(v):

458 + return v.ir_value() if hasattr(v, "ir_value") else v

459 +

460 + _init_iter = [

461 + _ival(running_max),

462 + _ival(running_sum),

463 + *[_ival(p) for p in acc_pv],

464 + ]

465 + _for_op = _scf.ForOp(

466 + c_zero_idx,

467 + trip_idx,

468 + c_one_idx,

469 + _init_iter,

470 + )

471 + _for_ip = ir.InsertionPoint(_for_op.body)

472 + _for_ip.__enter__()

473 + try:

474 + tg_idx = _for_op.induction_variable

475 + tg_i32 = fx.Int32(arith.index_cast(T.i32, tg_idx))

476 + partition_start = partition_base + tg_i32 * c_kcb

477 + bt_seq_base = seq * c_bt + (partition_start // fx.Int32(_BS))

478 + running_max = _for_op.inner_iter_args[0]

479 + running_sum = _for_op.inner_iter_args[1]

480 + acc_pv = list(_for_op.inner_iter_args[2:])

481 + for n_tile in range_constexpr(16):

482 + block_in_part = n_tile // _TILES_PER_BLOCK

483 + tile_in_block = n_tile % _TILES_PER_BLOCK

484 + tile_start_tok = partition_start + fx.Int32(n_tile * TILE_SIZE)

485 + tile_in_seq = tile_start_tok < seq_len

486 + bt_off = bt_seq_base + fx.Int32(block_in_part)

487 + bt_off_safe = tile_in_seq.select(bt_off, seq * c_bt)

488 + phys_block = buffer_ops.buffer_load(

489 + bt_rsrc,

490 + bt_off_safe,

491 + vec_width=1,

492 + dtype=T.i32,

493 + )

494 + block_base = phys_block * c_block

495 + data_region = block_base

496 + meta_region = block_base + c_meta_off

497 +

498 + # ``slot`` is the absolute slot index within the cache block

499 + # (range 0.._BS-1). For BS=16 it equals tok_in_tile; for BS=32

500 + # it equals tile_in_block * 16 + tok_in_tile. Used both for

501 + # the K/V data region and SoA metadata region addressing.

502 + slot = fx.Int32(tile_in_block * TILE_SIZE) + tok_in_tile

503 + data_bases_byte = (

504 + data_region

505 + + slot * (c_kv_heads * c_data_per_slot)

506 + + kv_h * c_data_per_slot

507 + )

508 + k_byte = data_bases_byte + chunk_in_tok * fx.Int32(16)

509 + k_packed = buffer_ops.buffer_load(

510 + kv_rsrc,

511 + k_byte // fx.Int32(4),

512 + vec_width=4,

513 + dtype=T.i32,

514 + )

515 +

516 + # ---- HOISTED: issue V data + meta HBM loads early ---------

517 + # These are async; their s_waitcnt is pushed by the compiler

518 + # past the K dequant + QK MFMA + softmax block, hiding most

519 + # of the V HBM latency behind compute.

520 + v_byte = data_bases_byte + c_keydata + chunk_in_tok * fx.Int32(16)

521 + v_packed = buffer_ops.buffer_load(

522 + kv_rsrc,

523 + v_byte // fx.Int32(4),

524 + vec_width=4,

525 + dtype=T.i32,

526 + )

527 + vscale_u16 = (

528 + meta_region // fx.Int32(2)

529 + + kv_h * c_meta_u16_per_kvh

530 + + c_vscale_off_u16

531 + + slot

532 + )

533 + vzero_u16 = (

534 + meta_region // fx.Int32(2)

535 + + kv_h * c_meta_u16_per_kvh

536 + + c_vzero_off_u16

537 + + slot

538 + )

539 + vscale_raw = buffer_ops.buffer_load(

540 + kv_rsrc,

541 + vscale_u16,

542 + vec_width=1,

543 + dtype=T.i16,

544 + )

545 + vzero_raw = buffer_ops.buffer_load(

546 + kv_rsrc,

547 + vzero_u16,

548 + vec_width=1,

549 + dtype=T.i16,

550 + )

551 +

552 + knorm_u16 = (

553 + meta_region // fx.Int32(2)

554 + + kv_h * c_meta_u16_per_kvh

555 + + c_knorm_off_u16

556 + + slot

557 + )

558 + knorm_raw = buffer_ops.buffer_load(

559 + kv_rsrc,

560 + knorm_u16,

561 + vec_width=1,

562 + dtype=T.i16,

563 + )

564 + knorm_f16 = arith.bitcast(T.f16, knorm_raw)

565 + knorm_f32 = arith.extf(T.f32, knorm_f16)

566 +

567 + # K dequant → LDS [token, head_dim] (natural)

568 + # Lane writes 32 bf16 (= 4×8) for token=tok_in_tile,

569 + # head_dims chunk_in_tok*32..+31 (4 sub-chunks of 8).

570 + tok_kreg = tok_in_tile * fx.Int32(HEAD_SIZE * 2 // 8)

571 + chunk_kreg = tok_kreg + chunk_in_tok * fx.Int32(8)

572 + for w in range_constexpr(4):

573 + word_i32 = vector.extract(k_packed, static_position=[w])

574 + bf16_elems = []

575 + for n in range_constexpr(8):

576 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

577 + nibble_idx = arith.index_cast(T.index, nibble)

578 + cent_f32 = cent_lds.load([nibble_idx])

579 + elem_bf16 = arith.trunc_f(T.bf16, cent_f32 * knorm_f32)

580 + bf16_elems.append(elem_bf16)

581 + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems)

582 + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16)

583 + vector.store(

584 + v_i64,

585 + kv_lds_i64,

586 + [arith.index_cast(T.index, chunk_kreg + fx.Int32(w * 2))],

587 + )

588 + gpu.barrier()

589 +

590 + # QK MFMA (CDNA4 wide-K): A=K[token, head_dim], B=Q (= Q^T).

591 + # K read: lane t = K[token=mfma_row, head_dim=chunk*32+col_grp*8..+7]

592 + # Same i64×2 indexing as Q.

593 + qk_acc = zero_v4

594 + for chk in range_constexpr(QK_K_CHUNKS):

595 + k_idx_i64 = (

596 + mfma_row * fx.Int32(HEAD_SIZE * 2 // 8)

597 + + fx.Int32(chk * 8)

598 + + mfma_col_grp * fx.Int32(2)

599 + )

600 + kv_load = vector.load_op(

601 + T.vec(2, T.i64),

602 + kv_lds_i64,

603 + [arith.index_cast(T.index, k_idx_i64)],

604 + )

605 + k_op = vector.bitcast(T.vec(8, T.bf16), kv_load)

606 + qk_acc = rocdl.mfma_f32_16x16x32_bf16(

607 + T.f32x4, [k_op, q_chunks[chk], qk_acc, 0, 0, 0]

608 + )

609 +

610 + # qk_acc layout: lane t holds C[token=(t/16)*4..+3, query=t%16]

611 + # 4 fp32/lane = 4 different tokens at SAME query_row.

612 +

613 + # Scale + mask out-of-context tokens.

614 + qk_acc = _vsplat_mul(qk_acc, QK_SCALE)

615 + for elem in range_constexpr(4):

616 + kv_tok = (

617 + partition_start

618 + + fx.Int32(n_tile * TILE_SIZE)

619 + + mfma_col_grp * fx.Int32(4)

620 + + fx.Int32(elem)

621 + )

622 + in_b = kv_tok < seq_len

623 + v = vector.extract(qk_acc, static_position=[elem])

624 + qk_acc = vector.insert(

625 + in_b.select(v, NEG_INF),

626 + qk_acc,

627 + static_position=[elem],

628 + dynamic_position=[],

629 + )

630 +

631 + # FA2 online softmax: per-query-row reduce.

632 + # Per-row max: max over 4 fp32 in lane, then xor-shuffle 16, 32

633 + # (across the 4 col_grps that share same mfma_row).

634 + local_max = vector.reduction(T.f32, "maxnumf", qk_acc)

635 + r1 = local_max.shuffle_xor(fx.Int32(16), c_w)

636 + local_max = local_max.maximumf(r1)

637 + r2 = local_max.shuffle_xor(fx.Int32(32), c_w)

638 + tile_max = local_max.maximumf(r2)

639 +

640 + new_max = running_max.maximumf(tile_max)

641 + max_diff = running_max - new_max

642 + safe_diff = (running_max > NEG_INF).select(max_diff, ZERO_F)

643 + scale = (safe_diff * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast)

644 + running_sum = running_sum * scale

645 + for h in range_constexpr(PV_N_CHUNKS):

646 + acc_pv[h] = _vsplat_mul(acc_pv[h], scale)

647 + running_max = new_max

648 +

649 + # Compute probs: p = exp((qk - new_max) * LOG2E)

650 + tile_sum = ZERO_F

651 + for elem in range_constexpr(4):

652 + s = vector.extract(qk_acc, static_position=[elem])

653 + d = s - new_max

654 + d = (new_max > NEG_INF).select(d, NEG_INF)

655 + p = (d * LOG2E_C).exp2(fastmath=arith.FastMathFlags.fast)

656 + tile_sum = tile_sum + p

657 + qk_acc = vector.insert(

658 + p, qk_acc, static_position=[elem], dynamic_position=[]

659 + )

660 +

661 + ts1 = tile_sum.shuffle_xor(fx.Int32(16), c_w)

662 + tile_sum = tile_sum + ts1

663 + ts2 = tile_sum.shuffle_xor(fx.Int32(32), c_w)

664 + tile_sum = tile_sum + ts2

665 + running_sum = running_sum + tile_sum

666 +

667 + # ---- V dequant → LDS (V/scale/zero raw loads were hoisted) -

668 + # v_packed, vscale_raw, vzero_raw are already in flight from

669 + # the top of this iteration; only the f32 conversion stays

670 + # here so the actual consume site is still close to the use.

671 + vscale_f32 = arith.extf(T.f32, arith.bitcast(T.f16, vscale_raw))

672 + vzero_f32 = arith.extf(T.f32, arith.bitcast(T.f16, vzero_raw))

673 +

674 + if USE_HW_TR:

675 + # V dequant → LDS [token][head_dim] (ROW-MAJOR, no transpose).

676 + # Each lane writes 32 contiguous bf16 (one token, head_dims

677 + # chunk_in_tok*32..+31) as 4× ds_write_b128 = 4× vec(2,i64).

678 + # Replaces 32× ds_write_b16 of the legacy transposed path.

679 + v_lds_elem_base = tok_in_tile * fx.Int32(

680 + HEAD_SIZE

681 + ) + chunk_in_tok * fx.Int32(32)

682 + for w in range_constexpr(4):

683 + word_i32 = vector.extract(v_packed, static_position=[w])

684 + bf16_elems = []

685 + for n in range_constexpr(8):

686 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

687 + nibble_f32 = arith.sitofp(T.f32, nibble)

688 + elem_f32 = nibble_f32 * vscale_f32 + vzero_f32

689 + elem_bf16 = arith.trunc_f(T.bf16, elem_f32)

690 + bf16_elems.append(elem_bf16)

691 + v_bf16 = vector.from_elements(T.vec(8, T.bf16), bf16_elems)

692 + v_i64 = vector.bitcast(T.vec(2, T.i64), v_bf16)

693 + v_lds_i64_idx = (v_lds_elem_base + fx.Int32(w * 8)) // fx.Int32(

694 + 4

695 + )

696 + vector.store(

697 + v_i64,

698 + kv_lds_i64,

699 + [arith.index_cast(T.index, v_lds_i64_idx)],

700 + )

701 + # ── HW V transpose: cross-lane LDS sync ─────────────────────

702 + # ds_read_tr16_b64 below introduces a cross-lane LDS read:

703 + # consumer lane t reads bytes written by source lanes

704 + # 0,4,8,12 (etc.). The compiler's automatic waitcnt insertion

705 + # is per-lane and has no awareness of the HW transpose's

706 + # cross-lane forwarding, so it can sink the next iteration's

707 + # ds_read ahead of these ds_write_b128 ops in the schedule.

708 + #

709 + # Two-part fence:

710 + # 1. sched_barrier(0): MachineScheduler reorder barrier

711 + # (mask=0 → no instruction class may cross). Pure hint,

712 + # generates no runtime instruction. Use the simple

713 + # sched_barrier (NOT sched_group_barrier, which is an

714 + # IGLP marker that requires partner barriers and

715 + # corrupts the schedule when used in isolation).

716 + # 2. s_waitcnt lgkmcnt=0: runtime drain of the LDS write

717 + # queue. vmcnt/expcnt left at no-wait so we don't

718 + # stall on speculatively-issued HBM loads (which can

719 + # include OOB-but-masked buffer_load reads from the

720 + # next K-tile's hoisted V prefetch).

721 + #

722 + # The SW path below does NOT need this because each lane

723 + # only reads cells it itself wrote (no cross-lane dep), and

724 + # the per-lane same-address waitcnt the compiler emits is

725 + # correct.

726 + #

727 + # Without this fence the reordered ds_read produces

728 + # wrong results on high-partition-count configs (large

729 + # padded num_partitions); low-partition configs are

730 + # unaffected.

731 + rocdl.sched_barrier(0)

732 + # encode_waitcnt(vmcnt=63, expcnt=7, lgkmcnt=0)

733 + # = 0xF | (7<<4) | (0<<8) | (3<<14) = 0xC07F

734 + rocdl.s_waitcnt(0xC07F)

735 + else:

736 + # Legacy transposed V_LDS path: 32 ds_write_b16 per lane.

737 + for w in range_constexpr(4):

738 + word_i32 = vector.extract(v_packed, static_position=[w])

739 + for n in range_constexpr(8):

740 + nibble = (word_i32 >> fx.Int32(n * 4)) & fx.Int32(0xF)

741 + nibble_f32 = arith.sitofp(T.f32, nibble)

742 + elem_f32 = nibble_f32 * vscale_f32 + vzero_f32

743 + elem_bf16 = arith.trunc_f(T.bf16, elem_f32)

744 + elem_i16 = arith.bitcast(T.i16, elem_bf16)

745 + head_dim = chunk_in_tok * fx.Int32(32) + fx.Int32(w * 8 + n)

746 + v_idx_i16 = head_dim * fx.Int32(TILE_SIZE) + tok_in_tile

747 + v_vec = vector.from_elements(T.vec(1, T.i16), [elem_i16])

748 + vector.store(

749 + v_vec,

750 + kv_lds_i16,

751 + [arith.index_cast(T.index, v_idx_i16)],

752 + )

753 + gpu.barrier()

754 +

755 + # ---- PV MFMA: A=V[head_dim, token], B=P (=qk_acc bf16) -----

756 + # P operand B layout matches qk_acc: lane t holds 4 bf16 at

757 + # K=token=(t/16)*4..+3, N=query=t%16. Just trunc_f to bf16.

758 + p_bf16 = arith.trunc_f(T.vec(4, T.bf16), qk_acc)

759 + p_op = vector.bitcast(T.vec(4, T.i16), p_bf16)

760 +

761 + if USE_HW_TR:

762 + # HW-transpose PV path: V_lds is row-major V[token][head_dim].

763 + # MFMA A operand layout: lane t holds 4 bf16 at

764 + # M=head_dim=mfma_row + h*16, K=token=(t/16)*4..+3

765 + # ds_read_tr16_b64 (4-element 16-bit transpose, per-16-lane block):

766 + # result[lane=t, elem=e] = Input[source_lane=e*4 + (t%16)//4,

767 + # col=t%4]

768 + # Per-lane address: token_idx = lane // 4 (covers 0..15 across

769 + # all four 16-lane MFMA blocks), hd_sub = (lane % 4)*4 selects

770 + # the 4-element column window inside the h*16 chunk.

771 + # Total LDS byte offset = kv_off + token_idx*HEAD_SIZE*2

772 + # + (h*16 + hd_sub)*2

773 + token_idx = lane >> fx.Int32(2)

774 + hd_sub = (lane & fx.Int32(3)) * fx.Int32(4)

775 + v_lane_byte = (

776 + fx.Int32(kv_off)

777 + + token_idx * fx.Int32(HEAD_SIZE * 2)

778 + + hd_sub * fx.Int32(2)

779 + )

780 + for h in range_constexpr(PV_N_CHUNKS):

781 + v_byte_off = v_lane_byte + fx.Int32(h * 32)

782 + v_byte_i64 = fx.Int64(v_byte_off)

783 + v_ptr = buffer_ops.create_llvm_ptr(

784 + v_byte_i64,

785 + address_space=3,

786 + )

787 + v_op_raw = rocdl.ds_read_tr16_b64(

788 + T.vec(4, T.i16),

789 + v_ptr,

790 + ).result

791 + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k(

792 + T.f32x4, [v_op_raw, p_op, acc_pv[h], 0, 0, 0]

793 + )

794 + else:

795 + # Legacy transposed-LDS PV: 1 ds_read_b64 per chunk.

796 + # Byte addr = (mfma_row + h*16) * TILE_SIZE * 2 + (col_grp*4) * 2

797 + # = (mfma_row + h*16) * 32 + col_grp*8

798 + # i64 idx = (mfma_row + h*16) * 4 + col_grp

799 + for h in range_constexpr(PV_N_CHUNKS):

800 + v_idx_i64 = (mfma_row + fx.Int32(h * 16)) * fx.Int32(

801 + 4

802 + ) + mfma_col_grp

803 + kv_load = vector.load_op(

804 + T.vec(1, T.i64),

805 + kv_lds_i64,

806 + [arith.index_cast(T.index, v_idx_i64)],

807 + )

808 + v_op = vector.bitcast(T.vec(4, T.i16), kv_load)

809 + acc_pv[h] = rocdl.mfma_f32_16x16x16bf16_1k(

810 + T.f32x4, [v_op, p_op, acc_pv[h], 0, 0, 0]

811 + )

812 +

813 + _scf.YieldOp(

814 + [

815 + _ival(running_max),

816 + _ival(running_sum),

817 + *[_ival(p) for p in acc_pv],

818 + ]

819 + )

820 + finally:

821 + _for_ip.__exit__(None, None, None)

822 + # Pull final accumulator values out of the for_op results.

823 + running_max = _for_op.results[0]

824 + running_sum = _for_op.results[1]

825 + acc_pv = list(_for_op.results[2:])

826 +

827 + # ===== STEP G: Output ===========================================

828 + # acc_pv[h] layout: lane t holds 4 fp32 at M=head_dim=(t/16)*4..+3 + h*16,

829 + # N=query_row=t%16.

830 + # 4 fp32/lane = 4 contiguous head_dim values for ONE query_row.

831 + safe_sum = (running_sum > ZERO_F).select(running_sum, ONE_F)

832 + rcp = ONE_F / safe_sum

833 +

834 + c_os = fx.Int32(_stride_out_seq)

835 + c_oh = fx.Int32(_stride_out_head)

836 + c_op_ = fx.Int32(_stride_out_part)

837 + out_base = seq * c_os + kv_h * c_oh + part * c_op_

838 +

839 + # Gate all global stores by mfma_row < QG. For QG=16 the predicate is

840 + # tautological (mfma_row ∈ [0,15]); for QG=8 lanes 8..15 (whose MFMA

841 + # outputs are computational waste) skip the stores entirely.

842 + valid_row_pred = arith.cmpi(

843 + arith.CmpIPredicate.ult,

844 + mfma_row.ir_value() if hasattr(mfma_row, "ir_value") else mfma_row,

845 + arith.constant(QG, type=T.i32),

846 + )

847 + _if = _scf.IfOp(valid_row_pred)

848 + with ir.InsertionPoint(_if.then_block):

849 + for h in range_constexpr(PV_N_CHUNKS):

850 + pv_norm = _vsplat_mul(acc_pv[h], rcp)

851 + pv_bf16 = arith.trunc_f(T.vec(4, T.bf16), pv_norm)

852 + pv_i32x2 = vector.bitcast(T.vec(2, T.i32), pv_bf16)

853 + head_dim_start = fx.Int32(h * 16) + mfma_col_grp * fx.Int32(4)

854 + out_off_elem = (

855 + out_base + mfma_row * fx.Int32(HEAD_SIZE) + head_dim_start

856 + )

857 + buffer_ops.buffer_store(

858 + pv_i32x2,

859 + out_rsrc,

860 + out_off_elem * fx.Int32(2),

861 + offset_is_bytes=True,

862 + )

863 +

864 + c_npq = fx.Int32(num_partitions * QG)

865 + ml_off = (

866 + seq * fx.Int32(_stride_ml_seq) + kv_h * c_npq + part * c_qg + mfma_row

867 + )

868 + es_off = (

869 + seq * fx.Int32(_stride_es_seq) + kv_h * c_npq + part * c_qg + mfma_row

870 + )

871 + buffer_ops.buffer_store(running_max, ml_rsrc, ml_off)

872 + buffer_ops.buffer_store(running_sum, es_rsrc, es_off)

873 + _scf.YieldOp([])

874 +

875 + return tq_decode_gqa6_kernel

vllm/v1/attention/ops/flydsl_turboquant_decode.py new

+874

@@ -0,0 +1,874 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""FlyDSL TurboQuant decode launcher (vLLM-side).

4 +

5 +Drop-in replacement for ``triton_turboquant_decode_attention_soa`` for the

6 +TQ decode profile HEAD_SIZE=128, MSE_BITS=4 K, VQB=4 V, N_CENTROIDS=16,

7 +BLOCK_SIZE in {16, 32}.

8 +

9 +Per-GQA kernel dispatch:

10 + * GQA group ∈ {8, 16} → canonical kernels.tq_decode (Qwen-class)

11 + * GQA group == 6 → kernels.tq_decode_gqa6 sibling (MiniMax-M2.5)

12 +

13 +Auto-selected on gfx950 when FlyDSL is importable; falls back to the SoA Triton

14 +decode otherwise (wrong arch, missing build tree). The GQA-6 sibling is imported

15 +best-effort: missing it does not affect Qwen GQA-{8,16} paths.

16 +

17 +Architecture:

18 + 1. Q rotation: ``q_rot = (query.float() @ PiT).bfloat16()`` — the same

19 + launcher-side rocBLAS GEMM the SoA Triton decode uses.

20 + 2. FlyDSL kernel writes per-partition outputs into

21 + ``[N, Hk, P, QG, D]`` bf16 + ``[N, Hk, P, QG]`` fp32 max/sum buffers.

22 + 3. A small Triton reducer combines partitions in the kernel's native layout

23 + (no permute/cast), writing the final ``[N, Hq, D]`` output.

24 +

25 +Kernel module is built once per ``(num_kv_heads, num_partitions,

26 +max_blocks_per_seq, scale)`` shape and cached.

27 +"""

28 +

29 +from __future__ import annotations

30 +

31 +from typing import Any

32 +

33 +import torch

34 +

35 +from vllm.logger import init_logger

36 +from vllm.triton_utils import tl, triton

37 +

38 +logger = init_logger(__name__)

39 +

40 +# -- FlyDSL import bootstrap --------------------------------------------------

41 +# FlyDSL is imported directly and must be importable (pip-installed, or on

42 +# PYTHONPATH). Import failures are caught below and fall back to the SoA Triton

43 +# decode, so vLLM startup is unaffected on hosts without FlyDSL.

44 +

45 +

46 +_FLYDSL_AVAILABLE: bool | None = None

47 +_TQ_MOD: Any = None # kernels.tq_decode module (Qwen GQA-{8,16})

48 +_TQ_MOD_GQA6: Any = None # kernels.tq_decode_gqa6 module (MiniMax GQA-6, optional)

49 +_FLYC: Any = None # flydsl.compiler

50 +_FX: Any = None # flydsl.expr

51 +_TYPING_T: Any = None

52 +_CC: Any = None # CompilationContext

53 +_IR: Any = None

54 +

55 +

56 +def is_flydsl_available() -> bool:

57 + """Return True iff on gfx950 and FlyDSL imports + kernel module load.

58 +

59 + The GQA-6 sibling kernel (``kernels.tq_decode_gqa6``) is imported

60 + best-effort: if it's missing (older FlyDSL checkout that pre-dates

61 + MiniMax support) the canonical Qwen path stays fully functional and

62 + only GQA-6 dispatches will fail with a clear error at launch time.

63 + """

64 + global _FLYDSL_AVAILABLE, _TQ_MOD, _TQ_MOD_GQA6

65 + global _FLYC, _FX, _TYPING_T, _CC, _IR

66 + if _FLYDSL_AVAILABLE is not None:

67 + return _FLYDSL_AVAILABLE

68 + try:

69 + # gfx950-only kernel (CDNA4 intrinsics); gate before importing FlyDSL.

70 + from vllm.platforms.rocm import on_gfx950

71 +

72 + if not on_gfx950():

73 + _FLYDSL_AVAILABLE = False

74 + return _FLYDSL_AVAILABLE

75 + import flydsl.compiler as flyc # noqa: F401

76 + import flydsl.expr as fx # noqa: F401

77 + from flydsl._mlir import ir

78 + from flydsl.compiler.kernel_function import CompilationContext

79 + from flydsl.expr.typing import T # noqa: F401

80 +

81 + # Kernel is vendored in-tree (ships with vLLM); FlyDSL provides only

82 + # the compiler/runtime framework (imported above).

83 + from vllm.v1.attention.ops.flydsl_kernels import tq_decode as tq_mod

84 +

85 + _FLYC = flyc

86 + _FX = fx

87 + _TYPING_T = T

88 + _CC = CompilationContext

89 + _IR = ir

90 + _TQ_MOD = tq_mod

91 + _FLYDSL_AVAILABLE = True

92 + logger.info_once("FlyDSL TQ decode launcher: available")

93 + except Exception as ex: # noqa: BLE001

94 + _FLYDSL_AVAILABLE = False

95 + logger.warning_once(

96 + "FlyDSL TQ decode launcher: unavailable (%s). "

97 + "Falling back to SoA Triton decode.",

98 + ex,

99 + )

100 + return _FLYDSL_AVAILABLE

101 + # Best-effort GQA-6 sibling import (does NOT gate Qwen availability).

102 + try:

103 + from vllm.v1.attention.ops.flydsl_kernels import (

104 + tq_decode_gqa6 as tq_mod_gqa6,

105 + )

106 +

107 + _TQ_MOD_GQA6 = tq_mod_gqa6

108 + logger.info_once("FlyDSL TQ decode GQA-6 sibling: available (MiniMax-class)")

109 + except Exception as ex: # noqa: BLE001

110 + _TQ_MOD_GQA6 = None

111 + logger.info_once(

112 + "FlyDSL TQ decode GQA-6 sibling: not available (%s); "

113 + "GQA-6 models will fall back to SoA Triton decode.",

114 + ex,

115 + )

116 + return _FLYDSL_AVAILABLE

117 +

118 +

119 +def is_flydsl_gqa6_available() -> bool:

120 + """True iff the optional GQA-6 sibling kernel module loaded.

121 +

122 + Used by the eligibility gate in turboquant_attn.py to decide whether

123 + a layer with num_kv_groups==6 can run on FlyDSL or must fall back

124 + to the SoA Triton decode.

125 + """

126 + if _FLYDSL_AVAILABLE is None:

127 + is_flydsl_available()

128 + return _TQ_MOD_GQA6 is not None

129 +

130 +

131 +# -- Kernel module cache -------------------------------------------------------

132 +# Compiling a FlyDSL kernel is expensive; cache by the constexprs that

133 +# parameterize stride math: num_kv_heads, num_partitions, max_blocks_per_seq,

134 +# scale. ``num_seqs`` is irrelevant to the kernel body so we omit it.

135 +_KERN_CACHE: dict[tuple, Any] = {}

136 +

137 +

138 +# -- One-shot log guards ------------------------------------------------------

139 +# Replaces vllm `logger.{info,warning}_once` (which hashes the format string

140 +# every call to dedup) with a true zero-overhead bool guard.

141 +_LOG_INVOKED_ONCE: bool = False

142 +_LOG_SINKS_WARNED: bool = False

143 +_LOG_NORM_WARNED: bool = False

144 +

145 +

146 +def _detect_max_capture_B() -> int:

147 + """Probe vLLM compilation config for the largest cudagraph capture size.

148 +

149 + Returns max(cudagraph_capture_sizes) if available, else falls back to

150 + 512 (vLLM's default cap). The segm-pool auto-grows if a larger batch is

151 + seen (see ``_SegmBufPool.get``).

152 +

153 + Used to size the segm-pool bucket once, so that all distinct B's

154 + (whether captured in a graph, or hit eagerly in mixed batches) share

155 + a single allocation per shape — eliminating the per-B growth that

156 + previously dominated dense-capture VRAM cost.

157 + """

158 + try:

159 + from vllm.config import get_current_vllm_config

160 +

161 + cfg = get_current_vllm_config()

162 + sizes = cfg.compilation_config.cudagraph_capture_sizes

163 + if sizes:

164 + return int(max(sizes))

165 + except Exception: # noqa: BLE001

166 + pass

167 + return 512

168 +

169 +

170 +class _SegmBufPool:

171 + """Single-bucket buffer pool for segm_out/segm_max/segm_sum/output.

172 +

173 + Each unique shape signature (Hk, Hq, num_partitions, QG, D, device,

174 + dtype) is backed by ONE allocation sized at ``max_B`` (= max captured

175 + cudagraph batch size, or env override). Per-call ``get(B, ...)``

176 + returns ``[:B]`` views — same data_ptr base, narrower N dim. The

177 + kernel grid is ``(B, ...)`` so it only ever writes the first B rows

178 + of segm_out / segm_max / segm_sum; downstream readers (Triton

179 + reducer + caller) only consume the first B rows of output.

180 +

181 + Why a single bucket: cudagraph records launch pointers at capture

182 + time. With per-B allocations, every captured size triggered a fresh

183 + allocation, so total VRAM scaled with the number of captured sizes.

184 + With single-bucket, the pool tops out at one max_B-sized buffer per

185 + shape, shared across ALL captured sizes and ALL eager mixed-batch

186 + calls. The data_ptr is stable from first allocation onward, so

187 + cudagraph capture is happy across all captured B's.

188 +

189 + Eliminates per-decode-step ``cudaMalloc`` (× 4) plus the device-side

190 + memset kernels behind ``torch.full(-inf)`` and ``torch.zeros``.

191 + The kernel always writes the FULL ``[B, Hk, P, QG, D]`` slice and

192 + every (n, kv_h, p, qg) position of segm_max / segm_sum (it stores

193 + ``-inf`` and ``0`` for empty partitions itself), so uninitialized

194 + bytes from the pool are safe to reuse for any B ≤ B_bucket.

195 + """

196 +

197 + __slots__ = ("_bufs", "_max_B")

198 +

199 + def __init__(self) -> None:

200 + self._bufs: dict[tuple, dict[str, torch.Tensor]] = {}

201 + self._max_B: int | None = None # lazy-init on first get()

202 +

203 + def get(

204 + self,

205 + B: int,

206 + Hk: int,

207 + Hq: int,

208 + num_partitions: int,

209 + QG: int,

210 + D: int,

211 + device: torch.device,

212 + q_dtype: torch.dtype,

213 + ) -> dict[str, torch.Tensor]:

214 + if self._max_B is None:

215 + self._max_B = _detect_max_capture_B()

216 + # B may exceed the detected max if user set max-num-seqs higher

217 + # than the largest captured size — still safe, we'll grow once.

218 + # Note: growing AFTER cudagraph capture would invalidate captured

219 + # pointers, so this should only happen during warmup or in

220 + # always-eager configs.

221 + B_bucket = max(self._max_B, int(B))

222 + key = (

223 + int(Hk),

224 + int(Hq),

225 + int(num_partitions),

226 + int(QG),

227 + int(D),

228 + str(device),

229 + q_dtype,

230 + )

231 + bufs = self._bufs.get(key)

232 + if bufs is None or bufs["segm_out"].shape[0] < B_bucket:

233 + if bufs is not None and bufs["segm_out"].shape[0] < B_bucket:

234 + # First-time grow: warn so user knows cudagraphs may be

235 + # invalidated. This should only happen during warmup or in

236 + # always-eager configs (a batch larger than any captured

237 + # cudagraph size).

238 + logger.warning_once(

239 + "FlyDSL _SegmBufPool: growing bucket from %d to %d "

240 + "(B=%d). If you see this AFTER cudagraph warmup, the "

241 + "previously-captured graphs hold stale pointers and "

242 + "will GPU-fault (max cudagraph capture size < B).",

243 + bufs["segm_out"].shape[0],

244 + B_bucket,

245 + B,

246 + )

247 + bufs = {

248 + "segm_out": torch.empty(

249 + (B_bucket, Hk, num_partitions, QG, D),

250 + dtype=torch.bfloat16,

251 + device=device,

252 + ),

253 + "segm_max": torch.empty(

254 + (B_bucket, Hk, num_partitions, QG),

255 + dtype=torch.float32,

256 + device=device,

257 + ),

258 + "segm_sum": torch.empty(

259 + (B_bucket, Hk, num_partitions, QG),

260 + dtype=torch.float32,

261 + device=device,

262 + ),

263 + "output": torch.empty(

264 + (B_bucket, Hq, D),

265 + dtype=q_dtype,

266 + device=device,

267 + ),

268 + # Stable pre-allocated buffer for the rotated query

269 + # tensor. On ROCm, fresh torch.empty / matmul-result allocations

270 + # after HIP graph capture can land in the graph memory pool and

271 + # cause GPU memory faults when passed as kernel arguments in the

272 + # eager mixed-batch path. By pooling q_rot here (same pattern as

273 + # segm_out/output above), the data_ptr is stable from first

274 + # allocation onward — same for all captured B's and eager calls.

275 + "q_rot": torch.empty(

276 + (B_bucket, Hq, D),

277 + dtype=q_dtype,

278 + device=device,

279 + ),

280 + # q_float: fp32 copy of query (bf16→fp32 for mm input).

281 + # Pooled to avoid fresh allocations post-capture.

282 + "q_float": torch.empty(

283 + (B_bucket, Hq, D),

284 + dtype=torch.float32,

285 + device=device,

286 + ),

287 + # q_rot_fp32: fp32 output of the rotation mm (before bf16 cast).

288 + # Using torch.mm(..., out=this) eliminates the fresh mm-result

289 + # allocation that would otherwise land in the HIP graph pool.

290 + "q_rot_fp32": torch.empty(

291 + (B_bucket, Hq, D),

292 + dtype=torch.float32,

293 + device=device,

294 + ),

295 + }

296 + self._bufs[key] = bufs

297 + self._max_B = B_bucket

298 + logger.info_once(

299 + "FlyDSL _SegmBufPool: allocated single-bucket "

300 + "shape=(Hk=%d, Hq=%d, P=%d, QG=%d, D=%d, dtype=%s) "

301 + "B_bucket=%d (covers all captured + eager B's). "

302 + "VRAM = %.1f MiB / shape.",

303 + Hk,

304 + Hq,

305 + num_partitions,

306 + QG,

307 + D,

308 + q_dtype,

309 + B_bucket,

310 + sum(t.numel() * t.element_size() for t in bufs.values()) / (1 << 20),

311 + )

312 + # Return [:B] views — same data_ptr base, narrower N dim.

313 + return {

314 + "segm_out": bufs["segm_out"][:B],

315 + "segm_max": bufs["segm_max"][:B],

316 + "segm_sum": bufs["segm_sum"][:B],

317 + "output": bufs["output"][:B],

318 + "q_rot": bufs["q_rot"][:B],

319 + "q_float": bufs["q_float"][:B],

320 + "q_rot_fp32": bufs["q_rot_fp32"][:B],

321 + }

322 +

323 + def stats(self) -> dict[str, int]:

324 + return {

325 + "shapes": len(self._bufs),

326 + "max_B": int(self._max_B or 0),

327 + "bytes": sum(

328 + sum(t.numel() * t.element_size() for t in d.values())

329 + for d in self._bufs.values()

330 + ),

331 + }

332 +

333 +

334 +_SEGM_POOL = _SegmBufPool()

335 +

336 +

337 +_HW_TR_CACHED: bool | None = None

338 +

339 +

340 +def _hw_tr_enabled() -> bool:

341 + """Resolve the HW V-transpose build flag.

342 +

343 + ON for gfx950+ (ds_read_tr16_b64 is bit-exact vs baseline and faster

344 + on Qwen-class shapes), off elsewhere. Resolved from the GPU arch.

345 + """

346 + global _HW_TR_CACHED

347 + if _HW_TR_CACHED is not None:

348 + return _HW_TR_CACHED

349 + try:

350 + from flydsl.runtime.device import get_rocm_arch as _arch

351 +

352 + a = str(_arch() or "")

353 + _HW_TR_CACHED = a.startswith("gfx950")

354 + except Exception: # noqa: BLE001

355 + _HW_TR_CACHED = False

356 + if _HW_TR_CACHED:

357 + logger.info_once("FlyDSL TQ: HW V transpose ON (default for gfx950+)")

358 + else:

359 + logger.info_once("FlyDSL TQ: HW V transpose OFF")

360 + return _HW_TR_CACHED

361 +

362 +

363 +_GET_KERNEL_STATS = {"hits": 0, "misses": 0, "build_total_s": 0.0}

364 +

365 +

366 +def _get_kernel(

367 + num_kv_heads: int,

368 + num_partitions: int,

369 + max_blocks_per_seq: int,

370 + scale: float,

371 + query_group_size: int,

372 + kv_block_size: int,

373 + use_hw_v_transpose: bool = False,

374 + num_seqs_hint: int = 1,

375 + tile_groups_per_partition: int = 1,

376 +):

377 + # ``num_seqs_hint`` (= runtime B) is forwarded to the build for shape

378 + # awareness; it does NOT participate in the cache key because the

379 + # kernel body uses gpu.block_idx.x (= seq index at runtime) and is

380 + # GQA-symmetric across B at build time.

381 + #

382 + # ``tile_groups_per_partition`` IS in the cache key — each

383 + # value compiles a different unrolled K-tile loop body.

384 + key = (

385 + num_kv_heads,

386 + int(num_partitions),

387 + int(max_blocks_per_seq),

388 + round(float(scale), 8),

389 + int(query_group_size),

390 + int(kv_block_size),

391 + bool(use_hw_v_transpose),

392 + int(tile_groups_per_partition),

393 + )

394 + cached = _KERN_CACHE.get(key)

395 + if cached is not None:

396 + _GET_KERNEL_STATS["hits"] += 1

397 + return cached

398 + assert is_flydsl_available()

399 + # Time the FlyDSL build path so we can quantify cudagraph capture cost.

400 + import time as _t

401 +

402 + _build_t0 = _t.perf_counter()

403 +

404 + # Per-GQA dispatch: GQA-6 (MiniMax-M2.5) lives in the sibling module

405 + # tq_decode_gqa6 to keep the Qwen kernel's invariants untouched.

406 + # GQA-{8,16} (Qwen) keep using the canonical tq_decode kernel.

407 + qg = int(query_group_size)

408 + if qg == 6:

409 + if _TQ_MOD_GQA6 is None:

410 + raise RuntimeError(

411 + "FlyDSL TQ GQA-6 sibling kernel "

412 + "(vllm.v1.attention.ops.flydsl_kernels.tq_decode_gqa6) "

413 + "failed to import; check that FlyDSL is installed/importable."

414 + )

415 + kmod = _TQ_MOD_GQA6

416 + kfn = kmod.build_tq_decode_gqa6_module(

417 + num_seqs=int(num_seqs_hint),

418 + num_kv_heads=num_kv_heads,

419 + num_partitions=num_partitions,

420 + max_blocks_per_seq=max_blocks_per_seq,

421 + softmax_scale=float(scale),

422 + query_group_size=qg,

423 + kv_block_size=int(kv_block_size),

424 + use_hw_v_transpose=bool(use_hw_v_transpose),

425 + tile_groups_per_partition=int(tile_groups_per_partition),

426 + )

427 + else:

428 + kmod = _TQ_MOD

429 + kfn = kmod.build_tq_decode_module(

430 + num_seqs=int(num_seqs_hint),

431 + num_kv_heads=num_kv_heads,

432 + num_partitions=num_partitions,

433 + max_blocks_per_seq=max_blocks_per_seq,

434 + softmax_scale=float(scale),

435 + query_group_size=qg,

436 + kv_block_size=int(kv_block_size),

437 + use_hw_v_transpose=bool(use_hw_v_transpose),

438 + tile_groups_per_partition=int(tile_groups_per_partition),

439 + )

440 + al = kmod.allocator

441 + block_threads = kmod.BLOCK_THREADS

442 +

443 + flyc = _FLYC

444 + fx = _FX

445 + T = _TYPING_T

446 + CompilationContext = _CC

447 + ir_mod = _IR

448 +

449 + @flyc.jit

450 + def _launch(

451 + out,

452 + es,

453 + ml,

454 + q,

455 + kvc,

456 + cents,

457 + bt,

458 + sl,

459 + gx: fx.Int32, # type: ignore[name-defined]

460 + gy: fx.Int32, # type: ignore[name-defined]

461 + gz: fx.Int32, # type: ignore[name-defined]

462 + stream: fx.Stream, # type: ignore[name-defined]

463 + ):

464 + # Re-finalize the LDS allocator for this launch (idempotent per build).

465 + al.finalized = False

466 + ctx = CompilationContext.get_current()

467 + with ir_mod.InsertionPoint(ctx.gpu_module_body):

468 + al.finalize()

469 + from flydsl.expr import arith

470 +

471 + grid_x = arith.index_cast(T.index, gx.ir_value())

472 + grid_y = arith.index_cast(T.index, gy.ir_value())

473 + grid_z = arith.index_cast(T.index, gz.ir_value())

474 + kfn(out, es, ml, q, kvc, cents, bt, sl).launch(

475 + grid=(grid_x, grid_y, grid_z),

476 + block=(block_threads, 1, 1),

477 + stream=stream,

478 + )

479 +

480 + _KERN_CACHE[key] = _launch

481 + _build_dt = _t.perf_counter() - _build_t0

482 + _GET_KERNEL_STATS["misses"] += 1

483 + _GET_KERNEL_STATS["build_total_s"] += _build_dt

484 + logger.info(

485 + "FlyDSL _get_kernel BUILD #%d dt=%.2fs (cumulative=%.1fs) key=%s",

486 + _GET_KERNEL_STATS["misses"],

487 + _build_dt,

488 + _GET_KERNEL_STATS["build_total_s"],

489 + key,

490 + )

491 + return _launch

492 +

493 +

494 +# -- Partition reducer (Triton, native kernel layout) -------------------------

495 +# Reduces the FA2 split-KV partials produced by the FlyDSL kernel.

496 +# segm_out [N, Hk, P, QG, D] bf16

497 +# segm_max [N, Hk, P, QG] fp32

498 +# segm_sum [N, Hk, P, QG] fp32

499 +# → output [N, Hq=Hk*QG, D] of any (bf16/fp16/fp32) dtype.

500 +@triton.jit

501 +def _reduce_partitions(

502 + output_ptr, # [N, Hq, D] in OUT_DTYPE

503 + segm_out_ptr, # [N, Hk, P, QG, D] bf16

504 + segm_max_ptr, # [N, Hk, P, QG] fp32

505 + segm_sum_ptr, # [N, Hk, P, QG] fp32

506 + out_stride_n: tl.int64, # stride on N axis (in elems of OUT_DTYPE)

507 + out_stride_h: tl.int64, # stride on Hq axis

508 + NUM_KV_HEADS: tl.constexpr,

509 + QG: tl.constexpr,

510 + NUM_PARTS: tl.constexpr,

511 + HEAD_SIZE: tl.constexpr,

512 +):

513 + # grid = (N, Hq)

514 + n = tl.program_id(0)

515 + hq = tl.program_id(1)

516 + kv_h = hq // QG

517 + qg = hq % QG

518 +

519 + # Per-partition pointers for this (n, kv_h, qg) row.

520 + # Strides in elems of source dtype.

521 + msum_base = n * (NUM_KV_HEADS * NUM_PARTS * QG) + kv_h * (NUM_PARTS * QG) + qg

522 + so_base = (

523 + n * (NUM_KV_HEADS * NUM_PARTS * QG * HEAD_SIZE)

524 + + kv_h * (NUM_PARTS * QG * HEAD_SIZE)

525 + + qg * HEAD_SIZE

526 + )

527 +

528 + p_off = tl.arange(0, NUM_PARTS)

529 + d_off = tl.arange(0, HEAD_SIZE)

530 +

531 + # Load per-partition max/sum.

532 + m_idx = msum_base + p_off * QG

533 + seg_max = tl.load(segm_max_ptr + m_idx)

534 + seg_sum = tl.load(segm_sum_ptr + m_idx)

535 +

536 + # Mask out empty partitions (max == -inf).

537 + valid = seg_max > float("-inf")

538 + overall_max = tl.max(tl.where(valid, seg_max, float("-inf")))

539 +

540 + # Rescale exp sums.

541 + rescale = tl.where(valid, tl.exp(seg_max - overall_max), 0.0)

542 + seg_sum_rescaled = seg_sum * rescale

543 + overall_sum = tl.sum(seg_sum_rescaled)

544 +

545 + # Load segment outputs and combine.

546 + so_idx = so_base + p_off[:, None] * (QG * HEAD_SIZE) + d_off[None, :]

547 + seg_out = tl.load(segm_out_ptr + so_idx).to(tl.float32)

548 + # acc = sum_p( seg_out_p * (seg_sum_p * rescale_p) ) / overall_sum

549 + weighted = seg_out * seg_sum_rescaled[:, None]

550 + acc_sum = tl.sum(weighted, axis=0)

551 + acc = tl.where(overall_sum > 0.0, acc_sum / overall_sum, 0.0)

552 +

553 + # Store.

554 + out_off = n * out_stride_n + hq * out_stride_h + d_off

555 + tl.store(output_ptr + out_off, acc.to(output_ptr.dtype.element_ty))

556 +

557 +

558 +# -- Public launcher ----------------------------------------------------------

559 +def flydsl_turboquant_decode_attention(

560 + query: torch.Tensor, # [B, Hq, D] bf16/fp16

561 + kv_cache: torch.Tensor, # [num_blocks, BS, Hk, slot_size_aligned]

562 + block_table: torch.Tensor, # [B, max_blocks_per_seq] int32

563 + seq_lens: torch.Tensor, # [B] int32

564 + Pi: torch.Tensor, # [D, D] fp32

565 + centroids: torch.Tensor, # [N_CENTROIDS] fp32

566 + scale: float,

567 + mse_bits: int,

568 + key_packed_size: int,

569 + value_quant_bits: int,

570 + value_packed_size: int,

571 + key_fp8: bool = False,

572 + norm_correction: bool = False,

573 + PiT: torch.Tensor | None = None,

574 + max_seq_len: int = 0,

575 + mid_o_buf: Any = None,

576 + output_buf: torch.Tensor | None = None,

577 + lse_buf: Any = None,

578 + buf_holder: Any = None,

579 + max_num_kv_splits: int = 32,

580 + sinks: torch.Tensor | None = None,

581 +) -> torch.Tensor:

582 + """SoA-decode-compatible launcher backed by the FlyDSL decode kernel.

583 +

584 + Constraints:

585 + * key_fp8 == False

586 + * mse_bits == 4

587 + * value_quant_bits == 4

588 + * centroids.numel() == 16

589 + * D == 128

590 + * block_size in {16, 32}

591 + * Hq // Hk in {6, 8, 16}

592 + - 8/16 → canonical tq_decode kernel (Qwen2.5-72B / Qwen3-32B)

593 + - 6 → tq_decode_gqa6 sibling kernel (MiniMax-M2.5)

594 +

595 + Sinks are NYI and silently ignored if set. norm_correction is honored

596 + implicitly via the pre-folded stored K-norm (see footer comment).

597 + """

598 + del mid_o_buf, lse_buf, key_packed_size, value_packed_size

599 + if not is_flydsl_available():

600 + raise RuntimeError(

601 + "FlyDSL decode requested but FlyDSL is not available (needs gfx950 "

602 + "+ importable FlyDSL); the SoA Triton decode is the fallback."

603 + )

604 + assert not key_fp8, "FlyDSL supports MSE-key path only"

605 + assert mse_bits == 4, f"FlyDSL expects mse_bits=4, got {mse_bits}"

606 + assert value_quant_bits == 4, (

607 + f"FlyDSL expects value_quant_bits=4, got {value_quant_bits}"

608 + )

609 +

610 + B, Hq, D = query.shape

611 + Hk = kv_cache.shape[2]

612 + block_size = kv_cache.shape[1]

613 + QG = Hq // Hk

614 + assert D == _TQ_MOD.HEAD_SIZE

615 + assert block_size in (16, 32), (

616 + f"FlyDSL supports kv_block_size 16 or 32, got {block_size}"

617 + )

618 + assert QG in (6, 8, 16), f"FlyDSL supports GQA factor 6, 8 or 16, got {QG}"

619 + if QG == 6 and _TQ_MOD_GQA6 is None:

620 + raise RuntimeError(

621 + "FlyDSL launcher: GQA-6 requested (MiniMax-class) but the "

622 + "tq_decode_gqa6 sibling module is not available. Update your "

623 + "FlyDSL checkout (must include kernels/tq_decode_gqa6.py); "

624 + "otherwise the SoA Triton decode is the fallback."

625 + )

626 + assert centroids.numel() == _TQ_MOD.N_CENTROIDS, (

627 + f"centroids.numel={centroids.numel()} != {_TQ_MOD.N_CENTROIDS}"

628 + )

629 +

630 + # ---- Per-layer cache for PiT_f32 + contiguous centroids ---

631 + # PiT and centroids are model constants (set once at layer warmup).

632 + # Avoid the per-decode-step transpose+cast and `.contiguous()` no-op

633 + # check by stashing the pre-cooked tensors on ``buf_holder`` (= layer).

634 + PiT_f32: torch.Tensor

635 + centroids_c: torch.Tensor

636 + if buf_holder is not None:

637 + PiT_f32 = getattr(buf_holder, "_tq_PiT_f32", None)

638 + if PiT_f32 is None:

639 + _PiT_src = PiT if PiT is not None else Pi.T.contiguous()

640 + PiT_f32 = (

641 + _PiT_src

642 + if _PiT_src.dtype == torch.float32

643 + else _PiT_src.to(torch.float32)

644 + )

645 + buf_holder._tq_PiT_f32 = PiT_f32

646 + centroids_c = getattr(buf_holder, "_tq_centroids_c", None)

647 + if centroids_c is None:

648 + centroids_c = centroids.contiguous()

649 + buf_holder._tq_centroids_c = centroids_c

650 + else:

651 + # Defensive fallback (unit tests may pass ``buf_holder=None``).

652 + _PiT_src = PiT if PiT is not None else Pi.T.contiguous()

653 + PiT_f32 = (

654 + _PiT_src if _PiT_src.dtype == torch.float32 else _PiT_src.to(torch.float32)

655 + )

656 + centroids_c = centroids.contiguous()

657 +

658 + # ---- Pooled q_rot (stable data_ptr post-capture) ------

659 + # On ROCm, fresh tensor allocations made AFTER HIP graph capture can land

660 + # in the graph memory pool and cause GPU memory faults when passed as kernel

661 + # pointer arguments in the eager mixed-batch path. We pool q_rot and the

662 + # fp32 intermediate the same way segm_out/output are pooled: one max-B

663 + # allocation per shape, reused across all B's.

664 + #

665 + # Note: pool_bufs is computed BELOW after num_partitions is determined;

666 + # q_rot needs pool_bufs["q_float"] and pool_bufs["q_rot"]. We compute the

667 + # rotation inline after the pool is fetched (see the pooled-buffers

668 + # section below).

669 +

670 + # ---- Partition count (FA2 split-KV) ----------------------------------

671 + #

672 + # FlyDSL runs INSIDE FULL cudagraph (TurboQuantMetadataBuilder._cudagraph_

673 + # support = UNIFORM_BATCH). That means the gridDim baked at capture

674 + # time MUST equal the gridDim at replay — the kernel launch parameters

675 + # are recorded into the captured graph. So `num_partitions` (which

676 + # becomes gridDim.z) MUST be derived from a stable source that produces

677 + # the same value at capture and at runtime.

678 + #

679 + # We use the worst-case context length implied by the block table's

680 + # allocation (block_table.shape[1] * block_size) as that source. It's

681 + # bounded by `max_model_len` (vLLM allocates the block table for the

682 + # configured max), and it's identical at capture and at runtime

683 + # because `block_table.shape[1]` is fixed once the model is loaded.

684 + # Per-tile OOB redirects + the masked FA-2 reducer ensure that for

685 + # short sequences the extra partitions contribute zero to the output.

686 + #

687 + # We always size from the worst case (block_table.shape[1] * block_size),

688 + # never from the per-step actual max_seq_len: a per-step size would make

689 + # the captured gridDim mismatch runtime and GPU-fault under cudagraph.

690 + kv_compute_block = _TQ_MOD.KV_COMPUTE_BLOCK

691 + worst_case_max_seq_len = int(block_table.shape[1]) * int(block_size)

692 + if max_seq_len <= 0:

693 + max_seq_len = worst_case_max_seq_len

694 + sizing_max_seq_len = worst_case_max_seq_len

695 +

696 + # ── Bounded num_partitions + internal tile-group looping ──

697 + # The kernel previously hardcoded one partition = 256 tokens (=

698 + # KV_COMPUTE_BLOCK), forcing num_partitions to scale linearly with

699 + # max_model_len, which made the cudagraph capture cost grow

700 + # prohibitively at long context. The new kernel takes

701 + # ``tile_groups_per_partition`` (TGPP); each CTA processes

702 + # ``TGPP * 16`` K-tiles = ``TGPP * KV_COMPUTE_BLOCK`` tokens with the

703 + # FA-2 online-softmax state accumulating across all of them. This

704 + # mirrors what the SoA Triton decode / HIP SoA-fusion already do.

705 + #

706 + # Strategy: cap num_partitions at ``MAX_PARTITIONS`` (default 32),

707 + # then derive TGPP so that ``num_partitions * TGPP * KV_COMPUTE_BLOCK

708 + # >= sizing_max_seq_len``.

709 + #

710 + # Examples (block_size=32 → worst_case max_bps*32):

711 + # max_model_len=8K: required=32, parts=32, TGPP=1 (no waste)

712 + # max_model_len=16K: required=64, parts=32, TGPP=2 (no waste)

713 + # max_model_len=32K: required=128, parts=32, TGPP=4 (no waste)

714 + # max_model_len=64K: required=256, parts=32, TGPP=8 (no waste)

715 + # max_model_len=128K: required=512, parts=32, TGPP=16

716 + MAX_PARTITIONS = 32

717 + required_num_partitions = (

718 + sizing_max_seq_len + kv_compute_block - 1

719 + ) // kv_compute_block

720 + # max_num_kv_splits acts as a parallelism floor (ensure at least this

721 + # many CTAs along grid.z), but is itself capped by MAX_PARTITIONS.

722 + parallelism_floor = min(MAX_PARTITIONS, max(1, max_num_kv_splits))

723 + # Cap num_partitions at MAX_PARTITIONS, but never go below the floor.

724 + num_partitions_actual = max(

725 + parallelism_floor,

726 + min(MAX_PARTITIONS, required_num_partitions),

727 + )

728 + # Round up to next power of 2 for the Triton reducer's tl.arange(0, N)

729 + # constraint (Triton requires power-of-2 ≥ 2).

730 + num_partitions = max(2, triton.next_power_of_2(num_partitions_actual))

731 + # Now derive tile_groups_per_partition (TGPP) so total coverage

732 + # ``num_partitions * TGPP * KV_COMPUTE_BLOCK`` is >= sizing_max_seq_len.

733 + # This guarantees no work is dropped at runtime regardless of seq_len.

734 + # Round TGPP up to the next power-of-2 so the JIT cache key is bounded

735 + # to a small set of values (e.g. {1, 2, 4, 8, 16}) — limits the number

736 + # of distinct kernel binaries that need to be compiled across requests.

737 + _tgpp_required = max(

738 + 1,

739 + (required_num_partitions + num_partitions - 1) // num_partitions,

740 + )

741 + tile_groups_per_partition = int(triton.next_power_of_2(_tgpp_required))

742 +

743 + # ---- Pooled buffers (no per-call cudaMalloc / memset) ----------

744 + # The FlyDSL kernel writes the FULL [B, Hk, P, QG, D] segm_out and the

745 + # FULL [B, Hk, P, QG] segm_max / segm_sum (running_max=-inf, sum=0 for

746 + # empty partitions are stored by the kernel itself), so uninitialized

747 + # buffers from the pool are safe to reuse.

748 + device = query.device

749 + pool_bufs = _SEGM_POOL.get(

750 + B,

751 + Hk,

752 + Hq,

753 + num_partitions,

754 + QG,

755 + D,

756 + device,

757 + query.dtype,

758 + )

759 + segm_out = pool_bufs["segm_out"]

760 + segm_max = pool_bufs["segm_max"]

761 + segm_sum = pool_bufs["segm_sum"]

762 + if output_buf is None:

763 + output = pool_bufs["output"]

764 + else:

765 + output = output_buf[:B] if output_buf.shape[0] != B else output_buf

766 +

767 + # ---- q rotation: GEMM q @ PiT via stable pooled buffers -------------

768 + _q_float = pool_bufs["q_float"] # [B, Hq, D] fp32, stable

769 + _q_rot_f32 = pool_bufs["q_rot_fp32"] # [B, Hq, D] fp32, stable mm output

770 + _q_rot_out = pool_bufs["q_rot"] # [B, Hq, D] query.dtype, stable

771 + _q_float.copy_(query) # bf16 → fp32 in-place, no alloc

772 + # mm into stable fp32 buffer via out= to avoid fresh allocations.

773 + # PiT_f32 is cached on buf_holder (stable ptr since first layer warmup).

774 + torch.mm(

775 + _q_float.view(B * Hq, D),

776 + PiT_f32,

777 + out=_q_rot_f32.view(B * Hq, D), # in-place into pool buffer

778 + )

779 + _q_rot_out.copy_(_q_rot_f32) # fp32 → bf16, into stable buf

780 + q_for_kernel = _q_rot_out # stable ptr for kernel launch

781 +

782 + # ---- FlyDSL kernel launch -------------------------------------------

783 + max_bps = int(block_table.shape[1])

784 + use_hw_tr = _hw_tr_enabled()

785 + launch = _get_kernel(

786 + Hk,

787 + num_partitions,

788 + max_bps,

789 + scale,

790 + QG,

791 + block_size,

792 + use_hw_v_transpose=use_hw_tr,

793 + num_seqs_hint=int(B),

794 + tile_groups_per_partition=int(tile_groups_per_partition),

795 + )

796 + # Zero-overhead one-shot info log (replaces logger.info_once which

797 + # hashes its format string on every call to dedup).

798 + global _LOG_INVOKED_ONCE, _LOG_SINKS_WARNED, _LOG_NORM_WARNED

799 + if not _LOG_INVOKED_ONCE:

800 + _LOG_INVOKED_ONCE = True

801 + logger.info(

802 + "FlyDSL launcher invoked (UNIFORM_BATCH cudagraph): "

803 + "B=%d Hk=%d Hq=%d D=%d QG=%d num_partitions=%d (actual=%d, "

804 + "cap=%d) TGPP=%d max_bps=%d block_size=%d max_seq_len=%d "

805 + "hw_v_transpose=%s (coverage=%d tokens, worst_case=%d tokens)",

806 + B,

807 + Hk,

808 + Hq,

809 + D,

810 + QG,

811 + num_partitions,

812 + num_partitions_actual,

813 + MAX_PARTITIONS,

814 + tile_groups_per_partition,

815 + max_bps,

816 + int(block_size),

817 + int(max_seq_len),

818 + use_hw_tr,

819 + num_partitions * tile_groups_per_partition * kv_compute_block,

820 + worst_case_max_seq_len,

821 + )

822 + launch(

823 + segm_out,

824 + segm_sum,

825 + segm_max,

826 + q_for_kernel,

827 + kv_cache,

828 + centroids_c,

829 + block_table,

830 + seq_lens,

831 + B,

832 + Hk,

833 + num_partitions,

834 + torch.cuda.current_stream(),

835 + )

836 +

837 + # ---- Reduce partitions -> [B, Hq, D] --------------------------------

838 + _reduce_partitions[(B, Hq)](

839 + output_ptr=output,

840 + segm_out_ptr=segm_out,

841 + segm_max_ptr=segm_max,

842 + segm_sum_ptr=segm_sum,

843 + out_stride_n=output.stride(0),

844 + out_stride_h=output.stride(1),

845 + NUM_KV_HEADS=Hk,

846 + QG=QG,

847 + NUM_PARTS=num_partitions,

848 + HEAD_SIZE=D,

849 + )

850 + if sinks is not None and not _LOG_SINKS_WARNED:

851 + _LOG_SINKS_WARNED = True

852 + logger.warning(

853 + "FlyDSL launcher: sinks ignored (NYI). Disable sinks if they "

854 + "are required for correctness."

855 + )

856 + # ── norm_correction is honored IMPLICITLY ───────────────────────────

857 + # When the model was stored with norm_correction=True (the *_nc presets

858 + # turboquant_4bit_nc / k3v4_nc / 3bit_nc), the per-token K-norm scalar

859 + # was pre-folded to ||k_t|| / ||c_t|| at store time by the TurboQuant

860 + # store path. The decode kernel just multiplies `c_vals * stored_knorm`,

861 + # which then equals `(c_vals / ||c_t||) * ||k_t||` — exactly the

862 + # unit-norm-renormalized centroid times the original key norm. The

863 + # Triton decode does the identical multiply, so both decoders honor

864 + # norm_correction equivalently and there is nothing to "do" at decode

865 + # time. The launcher's `norm_correction` arg is only kept for API

866 + # parity; we emit a one-shot INFO log to make the contract explicit.

867 + if norm_correction and not _LOG_NORM_WARNED:

868 + _LOG_NORM_WARNED = True

869 + logger.info(

870 + "FlyDSL launcher: norm_correction honored implicitly via "

871 + "pre-folded stored K-norm (cf. triton_turboquant_store step 3); "

872 + "no decode-time work required, identical to the SoA decode behavior."

873 + )

874 + return output

vllm/v1/attention/ops/turboquant_soa/__init__.py new

@@ -0,0 +1,3 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# HIP-free SoA Triton subset for the FlyDSL TurboQuant decode path.

vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode.py new

+643

@@ -0,0 +1,643 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Triton fused TurboQuant decode attention.

4 +

5 +Decode path: Triton stage1 (split-KV tiled attention scoring + value

6 +accumulation) + stage2 (log-sum-exp reduction across splits).

7 +

8 +Supports FP8 (E4M3) keys, 3-bit and 4-bit uniform quantized values.

9 +"""

10 +

11 +import math

12 +from typing import Any

13 +

14 +import torch

15 +

16 +from vllm.platforms import current_platform

17 +from vllm.triton_utils import tl, triton

18 +from vllm.v1.attention.ops.triton_decode_attention import (

19 + _fwd_kernel_stage2,

20 +)

21 +

22 +_FP8_E4B15: dict[int, int] = {}

23 +

24 +

25 +def _use_fp8_e4b15(device: int = 0) -> int:

26 + """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0.

27 + On non-CUDA platforms (e.g. XPU), always returns 0 (use e4nv format).

28 + """

29 + if device not in _FP8_E4B15:

30 + if current_platform.is_cuda_alike():

31 + cap = torch.cuda.get_device_capability(device)

32 + _FP8_E4B15[device] = 1 if cap < (8, 9) else 0

33 + else:

34 + _FP8_E4B15[device] = 0

35 + return _FP8_E4B15[device]

36 +

37 +

38 +# ---------------------------------------------------------------------------

39 +# Stage 1: Fused TQ score + value accumulation (BLOCK_KV tiled)

40 +# ---------------------------------------------------------------------------

41 +

42 +

43 +@triton.jit

44 +def _tq_decode_stage1(

45 + # Precomputed query projection

46 + Q_rot_ptr, # [B, Hq, D] float32

47 + # Compressed KV cache (combined K+V)

48 + KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8

49 + KV_cache_u16_ptr, # uint16 view of same storage — Opt#3 SoA loads

50 + # Block table and sequence info

51 + Block_table_ptr, # [B, max_num_blocks] int32

52 + Seq_lens_ptr, # [B] int32

53 + # TQ parameters

54 + Centroids_ptr, # [n_centroids] float32

55 + # Output (intermediate for stage2)

56 + Mid_o_ptr, # [B, Hq, NUM_KV_SPLITS, D+1] float32

57 + # Sink support

58 + Sink_ptr, # [Hq] float32, pre-computed sink attention logits per head

59 + # Strides

60 + stride_qb,

61 + stride_qh, # Q strides: [B, Hq, D]

62 + stride_cache_block: tl.int64, # bytes per block (bs*H*slot_aligned)

63 + stride_bt_b: tl.int64, # block_table stride per batch

64 + stride_mid_b,

65 + stride_mid_h,

66 + stride_mid_s, # mid_o strides

67 + # Constexpr dims

68 + NUM_KV_HEADS: tl.constexpr,

69 + HEAD_DIM: tl.constexpr,

70 + BLOCK_SIZE: tl.constexpr, # KV cache block_size (pages)

71 + NUM_KV_SPLITS: tl.constexpr,

72 + KV_GROUP_SIZE: tl.constexpr, # Hq // Hk

73 + # TQ layout constants

74 + MSE_BITS: tl.constexpr, # 3 or 4

75 + MSE_BYTES: tl.constexpr, # ceil(D * mse_bits / 8)

76 + VQB: tl.constexpr, # value_quant_bits (4 or 8=FP8)

77 + VAL_DATA_BYTES: tl.constexpr, # ceil(D * vqb / 8) or D for FP8

78 + # Opt#3 SoA layout

79 + KEY_DATA_BYTES: tl.constexpr, # MSE_BYTES (MSE) or HEAD_DIM (FP8)

80 + META_REGION_OFFSET: tl.constexpr, # bytes from block start to SoA region

81 + NUM_SOA_FIELDS: tl.constexpr, # 3 (MSE) or 2 (FP8)

82 + SOA_K_NORM: tl.constexpr, # 0 (MSE); unused for FP8

83 + SOA_V_SCALE: tl.constexpr, # 1 (MSE) / 0 (FP8)

84 + SOA_V_ZERO: tl.constexpr, # 2 (MSE) / 1 (FP8)

85 + # Score constants

86 + ATTN_SCALE: tl.constexpr, # 1/sqrt(D)

87 + # Block tile sizes

88 + BLOCK_D: tl.constexpr, # next_power_of_2(HEAD_DIM)

89 + BLOCK_KV: tl.constexpr, # tokens per tile (16)

90 + KEY_FP8: tl.constexpr, # 1 if K is stored as FP8

91 + NORM_CORRECTION: tl.constexpr = 0, # 1 = re-normalize centroids

92 + FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)

93 + USE_SINKS: tl.constexpr = 0, # 1 = use sink tokens for attention anchoring

94 +):

95 + bid = tl.program_id(0) # batch index

96 + hid = tl.program_id(1) # q_head index

97 + sid = tl.program_id(2) # kv_split index

98 +

99 + kv_head = hid // KV_GROUP_SIZE

100 +

101 + # Sequence length for this batch

102 + seq_len = tl.load(Seq_lens_ptr + bid)

103 +

104 + # KV split range

105 + split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)

106 + split_start = split_len * sid

107 + split_end = tl.minimum(split_start + split_len, seq_len)

108 +

109 + if split_start >= split_end:

110 + return

111 +

112 + # Dimension offsets

113 + d_offs = tl.arange(0, BLOCK_D)

114 + d_mask = d_offs < HEAD_DIM

115 + kv_range = tl.arange(0, BLOCK_KV)

116 +

117 + # Load query vector: q_rot — [BLOCK_D] float32

118 + q_base = bid * stride_qb + hid * stride_qh

119 + q_rot = tl.load(Q_rot_ptr + q_base + d_offs, mask=d_mask, other=0.0).to(tl.float32)

120 +

121 + # Precompute byte/bit index vectors for MSE gather loads

122 + if not KEY_FP8:

123 + mse_bit_off = d_offs * MSE_BITS

124 + mse_byte_idx = mse_bit_off // 8

125 + mse_bit_shift = mse_bit_off % 8

126 + mse_mask = (1 << MSE_BITS) - 1

127 +

128 + # Precompute value bit/byte index vectors (loop-invariant)

129 + if VQB == 3:

130 + val_bit_off = d_offs * 3

131 + val_byte_idx = val_bit_off // 8

132 + val_bit_shift = val_bit_off % 8

133 +

134 + # Online softmax accumulators

135 + # Sink tokens provide a pre-computed attention bias that should be

136 + # included in the softmax normalization. For the first split (sid==0),

137 + # we initialize m_prev with the sink logit. L is always 1.0 for all

138 + # splits to match unified attention semantics.

139 + if USE_SINKS:

140 + if sid == 0:

141 + m_prev = tl.load(Sink_ptr + hid).to(tl.float32)

142 + l_prev = 1.0

143 + else:

144 + m_prev = -float("inf")

145 + l_prev = 0.0

146 + else:

147 + m_prev = -float("inf")

148 + l_prev = 0.0

149 + acc = tl.zeros([BLOCK_D], dtype=tl.float32)

150 +

151 + bt_base = bid * stride_bt_b

152 +

153 + # ================================================================

154 + # TILED LOOP: process BLOCK_KV tokens per iteration

155 + # ================================================================

156 + for start_n in range(split_start, split_end, BLOCK_KV):

157 + kv_offs = start_n + kv_range

158 + kv_mask = kv_offs < split_end

159 +

160 + page_idx = kv_offs // BLOCK_SIZE

161 + page_off = kv_offs % BLOCK_SIZE

162 + block_nums = tl.load(

163 + Block_table_ptr + bt_base + page_idx,

164 + mask=kv_mask,

165 + other=0,

166 + ).to(tl.int64)

167 +

168 + # Opt#3 SoA addressing: data region then metadata region per block.

169 + slot_within_block = page_off.to(tl.int64)

170 + block_base = block_nums * stride_cache_block

171 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

172 + data_bases = (

173 + block_base

174 + + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)

175 + + tl.cast(kv_head, tl.int64) * DATA_BYTES_PER_SLOT

176 + )

177 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(

178 + kv_head, tl.int64

179 + ) * (NUM_SOA_FIELDS * BLOCK_SIZE)

180 + knorm_u16_addrs = (

181 + head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block

182 + )

183 + vscale_u16_addrs = (

184 + head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block

185 + )

186 + vzero_u16_addrs = (

187 + head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block

188 + )

189 +

190 + # ============================================================

191 + # COMPUTE ATTENTION SCORES: [BLOCK_KV]

192 + # ============================================================

193 + if KEY_FP8:

194 + k_addrs = data_bases[:, None] + d_offs[None, :]

195 + k_raw = tl.load(

196 + KV_cache_ptr + k_addrs,

197 + mask=kv_mask[:, None] & d_mask[None, :],

198 + other=0,

199 + )

200 + if FP8_E4B15:

201 + k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)

202 + else:

203 + k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)

204 + scores = (

205 + tl.sum(

206 + tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0),

207 + axis=1,

208 + )

209 + * ATTN_SCALE

210 + )

211 + scores = tl.where(kv_mask, scores, -float("inf"))

212 + else:

213 + # MSE unpack + norms

214 + mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]

215 + mse_raw0 = tl.load(

216 + KV_cache_ptr + mse_addrs0,

217 + mask=kv_mask[:, None] & d_mask[None, :],

218 + other=0,

219 + ).to(tl.int32)

220 + mse_raw1 = tl.load(

221 + KV_cache_ptr + mse_addrs0 + 1,

222 + mask=kv_mask[:, None] & d_mask[None, :],

223 + other=0,

224 + ).to(tl.int32)

225 + raw16 = mse_raw0 | (mse_raw1 << 8)

226 + mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask

227 +

228 + # Centroid gather + dot product

229 + c_vals = tl.load(

230 + Centroids_ptr + mse_idx,

231 + mask=kv_mask[:, None] & d_mask[None, :],

232 + other=0.0,

233 + )

234 +

235 + # Opt#1: norm-correction (1/||c_vec||) is pre-folded into the

236 + # stored per-token scalar by the MSE store kernel when

237 + # NORM_CORRECTION=1, so we skip the per-tile sum+sqrt+divide.

238 + # NORM_CORRECTION is kept in the signature for API parity; it is

239 + # implicit in the stored norm value now.

240 +

241 + term1 = tl.sum(

242 + tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0),

243 + axis=1,

244 + )

245 +

246 + # Opt#3: K-norms live in the per-block SoA region. Single u16

247 + # load per token; contiguous across tokens within a block.

248 + norm_u16 = tl.load(

249 + KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask, other=0

250 + )

251 + vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)

252 +

253 + scores = vec_norms * term1 * ATTN_SCALE

254 + scores = tl.where(kv_mask, scores, -float("inf"))

255 +

256 + # ============================================================

257 + # ONLINE SOFTMAX UPDATE (block-level)

258 + # ============================================================

259 + n_e_max = tl.maximum(tl.max(scores, 0), m_prev)

260 + re_scale = tl.exp(m_prev - n_e_max)

261 + p = tl.exp(scores - n_e_max)

262 +

263 + # ============================================================

264 + # VALUE LOAD + DEQUANTIZE: [BLOCK_KV, BLOCK_D]

265 + # Opt#3: V data immediately follows K data in the slot (no

266 + # interleaved metadata). Scale/zero come from the SoA metadata

267 + # region as single u16 loads per field per token.

268 + # ============================================================

269 + val_bases = data_bases + KEY_DATA_BYTES

270 +

271 + if VQB == 3:

272 + val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]

273 + val_raw0 = tl.load(

274 + KV_cache_ptr + val_addrs0,

275 + mask=kv_mask[:, None] & d_mask[None, :],

276 + other=0,

277 + ).to(tl.int32)

278 + val_raw1 = tl.load(

279 + KV_cache_ptr + val_addrs0 + 1,

280 + mask=kv_mask[:, None] & d_mask[None, :],

281 + other=0,

282 + ).to(tl.int32)

283 + raw16 = val_raw0 | (val_raw1 << 8)

284 + v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)

285 + else: # VQB == 4

286 + vb_idx = d_offs // 2

287 + vb_shift = (d_offs % 2) * 4

288 + val_addrs = val_bases[:, None] + vb_idx[None, :]

289 + val_raw = tl.load(

290 + KV_cache_ptr + val_addrs,

291 + mask=kv_mask[:, None] & d_mask[None, :],

292 + other=0,

293 + ).to(tl.int32)

294 + v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)

295 +

296 + scale_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask, other=0)

297 + zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask, other=0)

298 + v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)

299 + v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)

300 + values = v_idx * v_scales[:, None] + v_zeros[:, None]

301 +

302 + # ============================================================

303 + # WEIGHTED VALUE ACCUMULATION

304 + # ============================================================

305 + acc = acc * re_scale + tl.sum(p[:, None] * values, 0)

306 + l_prev = l_prev * re_scale + tl.sum(p, 0)

307 + m_prev = n_e_max

308 +

309 + # Store partial result

310 + out_base = bid * stride_mid_b + hid * stride_mid_h + sid * stride_mid_s

311 + safe_l = tl.where(l_prev > 0.0, l_prev, 1.0)

312 + tl.store(Mid_o_ptr + out_base + d_offs, acc / safe_l, mask=d_mask)

313 + lse = m_prev + tl.log(safe_l)

314 + tl.store(Mid_o_ptr + out_base + HEAD_DIM, lse)

315 +

316 +

317 +# ---------------------------------------------------------------------------

318 +# Pre-dequant kernel: Bulk dequant K (MSE+norms) and V to fp16

319 +# ---------------------------------------------------------------------------

320 +

321 +

322 +@triton.jit

323 +def _tq_full_dequant_kv(

324 + KV_cache_ptr,

325 + KV_cache_u16_ptr,

326 + Block_table_ptr,

327 + Centroids_ptr,

328 + K_out_ptr, # [B, Hk, max_seq, D] float16

329 + V_out_ptr, # [B, Hk, max_seq, D] float16

330 + stride_ko_b: tl.int64,

331 + stride_ko_h: tl.int64,

332 + stride_ko_s: tl.int64,

333 + stride_vo_b: tl.int64,

334 + stride_vo_h: tl.int64,

335 + stride_vo_s: tl.int64,

336 + stride_cache_block: tl.int64,

337 + stride_bt_b: tl.int64,

338 + HEAD_DIM: tl.constexpr,

339 + BLOCK_SIZE: tl.constexpr,

340 + NUM_KV_HEADS: tl.constexpr,

341 + MSE_BYTES: tl.constexpr,

342 + VQB: tl.constexpr,

343 + VAL_DATA_BYTES: tl.constexpr,

344 + MSE_BITS: tl.constexpr,

345 + KEY_FP8: tl.constexpr,

346 + # Opt#3 SoA layout

347 + KEY_DATA_BYTES: tl.constexpr,

348 + META_REGION_OFFSET: tl.constexpr,

349 + NUM_SOA_FIELDS: tl.constexpr,

350 + SOA_K_NORM: tl.constexpr,

351 + SOA_V_SCALE: tl.constexpr,

352 + SOA_V_ZERO: tl.constexpr,

353 + BLOCK_D: tl.constexpr,

354 + NORM_CORRECTION: tl.constexpr = 0,

355 + FP8_E4B15: tl.constexpr = 0, # 1 = use e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)

356 +):

357 + """Full dequant: reconstruct K (MSE centroids * norm or FP8) and V to fp16.

358 +

359 + Opt#3 SoA layout: data region [bs, H, KD+VD] + per-block SoA metadata.

360 + """

361 + pos = tl.program_id(0)

362 + bh = tl.program_id(1)

363 + bid = bh // NUM_KV_HEADS

364 + hid = bh % NUM_KV_HEADS

365 +

366 + page_idx = pos // BLOCK_SIZE

367 + page_off = pos % BLOCK_SIZE

368 + block_num = tl.load(Block_table_ptr + bid * stride_bt_b + page_idx).to(tl.int64)

369 +

370 + block_base = block_num * stride_cache_block

371 + slot_within_block = tl.cast(page_off, tl.int64)

372 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

373 + data_base = (

374 + block_base

375 + + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)

376 + + tl.cast(hid, tl.int64) * DATA_BYTES_PER_SLOT

377 + )

378 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(

379 + hid, tl.int64

380 + ) * (NUM_SOA_FIELDS * BLOCK_SIZE)

381 + knorm_u16_addr = head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block

382 + vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block

383 + vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block

384 +

385 + d_offs = tl.arange(0, BLOCK_D)

386 + d_mask = d_offs < HEAD_DIM

387 +

388 + # === K dequant ===

389 + ko_base = bid * stride_ko_b + hid * stride_ko_h + pos * stride_ko_s

390 + if KEY_FP8:

391 + k_raw = tl.load(KV_cache_ptr + data_base + d_offs, mask=d_mask, other=0)

392 + if FP8_E4B15:

393 + k_recon = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)

394 + else:

395 + k_recon = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)

396 + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)

397 + else:

398 + mse_bit_off = d_offs * MSE_BITS

399 + mse_byte_idx = mse_bit_off // 8

400 + mse_bit_shift = mse_bit_off % 8

401 + mse_umask = (1 << MSE_BITS) - 1

402 +

403 + mse_raw0 = tl.load(

404 + KV_cache_ptr + data_base + mse_byte_idx, mask=d_mask, other=0

405 + ).to(tl.int32)

406 + mse_raw1 = tl.load(

407 + KV_cache_ptr + data_base + mse_byte_idx + 1, mask=d_mask, other=0

408 + ).to(tl.int32)

409 + raw16_key = mse_raw0 | (mse_raw1 << 8)

410 + mse_idx = (raw16_key >> mse_bit_shift) & mse_umask

411 +

412 + k_mse = tl.load(Centroids_ptr + mse_idx, mask=d_mask, other=0.0)

413 +

414 + # Opt#1: norm-correction is pre-folded into the stored scalar.

415 + # Opt#3: K-norm lives in the SoA region — single u16 load.

416 + norm_u16 = tl.load(KV_cache_u16_ptr + knorm_u16_addr)

417 + vec_norm = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)

418 +

419 + k_recon = vec_norm * k_mse

420 + tl.store(K_out_ptr + ko_base + d_offs, k_recon.to(tl.float16), mask=d_mask)

421 +

422 + # === V dequant ===

423 + val_base = data_base + KEY_DATA_BYTES

424 + if VQB == 4:

425 + vb_idx = d_offs // 2

426 + vb_shift = (d_offs % 2) * 4

427 + val_raw = tl.load(KV_cache_ptr + val_base + vb_idx, mask=d_mask, other=0).to(

428 + tl.int32

429 + )

430 + v_idx = ((val_raw >> vb_shift) & 0xF).to(tl.float32)

431 + elif VQB == 3:

432 + val_bit_off = d_offs * 3

433 + val_byte_idx = val_bit_off // 8

434 + val_bit_shift = val_bit_off % 8

435 + val_raw0 = tl.load(

436 + KV_cache_ptr + val_base + val_byte_idx, mask=d_mask, other=0

437 + ).to(tl.int32)

438 + val_raw1 = tl.load(

439 + KV_cache_ptr + val_base + val_byte_idx + 1, mask=d_mask, other=0

440 + ).to(tl.int32)

441 + raw16_val = val_raw0 | (val_raw1 << 8)

442 + v_idx = ((raw16_val >> val_bit_shift) & 0x7).to(tl.float32)

443 + else:

444 + v_idx = tl.zeros([BLOCK_D], dtype=tl.float32)

445 +

446 + # Opt#3: V scale/zero single-u16 loads from SoA region.

447 + sc_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addr)

448 + zr_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addr)

449 + v_scale = sc_u16.to(tl.float16, bitcast=True).to(tl.float32)

450 + v_zero = zr_u16.to(tl.float16, bitcast=True).to(tl.float32)

451 + v_vals = v_idx * v_scale + v_zero

452 +

453 + vo_base = bid * stride_vo_b + hid * stride_vo_h + pos * stride_vo_s

454 + tl.store(V_out_ptr + vo_base + d_offs, v_vals.to(tl.float16), mask=d_mask)

455 +

456 +

457 +# ---------------------------------------------------------------------------

458 +# Stage 2: Reuse from triton_decode_attention.py

459 +# ---------------------------------------------------------------------------

460 +

461 +# ---------------------------------------------------------------------------

462 +# Launcher — cached constants + fused GEMM

463 +# ---------------------------------------------------------------------------

464 +

465 +_layout_cache: dict = {}

466 +

467 +

468 +def _get_layout(D, mse_bits, value_quant_bits, key_packed_size):

469 + """Get cached layout constants."""

470 + key = (D, mse_bits, value_quant_bits, key_packed_size)

471 + cfg = _layout_cache.get(key)

472 + if cfg is None:

473 + val_data_bytes = math.ceil(D * value_quant_bits / 8)

474 + cfg = {

475 + "mse_bytes": math.ceil(D * mse_bits / 8),

476 + "val_data_bytes": val_data_bytes,

477 + "mse_bits": mse_bits,

478 + "n_centroids": 2**mse_bits,

479 + "BLOCK_D": triton.next_power_of_2(D),

480 + }

481 + _layout_cache[key] = cfg

482 + return cfg

483 +

484 +

485 +def triton_turboquant_decode_attention(

486 + query: torch.Tensor, # [B, Hq, D] — original query

487 + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8

488 + block_table: torch.Tensor, # [B, max_num_blocks] int32

489 + seq_lens: torch.Tensor, # [B] int32

490 + Pi: torch.Tensor, # [D, D] float32

491 + centroids: torch.Tensor, # [n_centroids] float32

492 + scale: float,

493 + mse_bits: int,

494 + key_packed_size: int,

495 + value_quant_bits: int,

496 + key_fp8: bool = False,

497 + norm_correction: bool = False,

498 + PiT: torch.Tensor | None = None, # [D, D] pre-computed Pi.T contiguous

499 + # Pre-allocated buffers (optional, avoids per-call allocation)

500 + mid_o_buf: torch.Tensor | None = None,

501 + output_buf: torch.Tensor | None = None,

502 + lse_buf: torch.Tensor | None = None,

503 + buf_holder: Any = None,

504 + max_num_kv_splits: int = 32, # fixed split count (must be constant for cudagraph)

505 + sinks: torch.Tensor | None = None, # [Hq] float32, pre-computed sink logits

506 +) -> torch.Tensor:

507 + """Launch fused TQ decode attention (Triton stage1 + stage2).

508 +

509 + Returns: output tensor [B, Hq, D] in query's dtype.

510 + """

511 + B, Hq, D = query.shape

512 + Hk = kv_cache.shape[2]

513 + block_size = kv_cache.shape[1]

514 + kv_group_size = Hq // Hk

515 + device = query.device

516 +

517 + cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)

518 +

519 + # Opt#3 SoA layout constants (derived locally; match store-side values).

520 + key_data_bytes = D if key_fp8 else cfg["mse_bytes"]

521 + data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"]

522 + meta_region_offset = block_size * Hk * data_bytes_per_slot

523 + num_soa_fields = 2 if key_fp8 else 3

524 + soa_k_norm = 0

525 + soa_v_scale = 0 if key_fp8 else 1

526 + soa_v_zero = 1 if key_fp8 else 2

527 + kv_cache_u16 = kv_cache.view(torch.uint16)

528 +

529 + # Compute q_rot = q @ Pi.T (rotated query for MSE key scoring)

530 + # FP8 path: pass query directly (float16); kernel casts inline.

531 + # MSE path: still needs external GEMM (cuBLAS), so q_rot is float32.

532 + if key_fp8:

533 + q_rot = query.contiguous()

534 + else:

535 + q_float = query.float()

536 + if PiT is None:

537 + PiT = Pi.T.contiguous()

538 + q_rot = (q_float @ PiT).contiguous()

539 +

540 + NUM_KV_SPLITS = max_num_kv_splits

541 +

542 + if (

543 + mid_o_buf is not None

544 + and mid_o_buf.shape[0] >= B

545 + and mid_o_buf.shape[2] >= NUM_KV_SPLITS

546 + ):

547 + mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]

548 + else:

549 + mid_o = torch.empty(

550 + B,

551 + Hq,

552 + NUM_KV_SPLITS,

553 + D + 1,

554 + dtype=torch.float32,

555 + device=device,

556 + )

557 + if buf_holder is not None:

558 + buf_holder._tq_mid_o_buf = mid_o

559 +

560 + # Stage 1: split-KV tiled attention scoring + value accumulation

561 + fp8_e4b15 = _use_fp8_e4b15(device.index or 0)

562 + BLOCK_KV = 4

563 + grid = (B, Hq, NUM_KV_SPLITS)

564 +

565 + # Prepare sink pointer (convert to float32 on device if provided)

566 + use_sinks = sinks is not None

567 +

568 + _tq_decode_stage1[grid](

569 + q_rot,

570 + kv_cache,

571 + kv_cache_u16,

572 + block_table,

573 + seq_lens,

574 + centroids,

575 + mid_o,

576 + sinks,

577 + q_rot.stride(0),

578 + q_rot.stride(1),

579 + kv_cache.stride(0),

580 + block_table.stride(0),

581 + mid_o.stride(0),

582 + mid_o.stride(1),

583 + mid_o.stride(2),

584 + NUM_KV_HEADS=Hk,

585 + HEAD_DIM=D,

586 + BLOCK_SIZE=block_size,

587 + NUM_KV_SPLITS=NUM_KV_SPLITS,

588 + KV_GROUP_SIZE=kv_group_size,

589 + MSE_BITS=mse_bits,

590 + MSE_BYTES=cfg["mse_bytes"],

591 + VQB=value_quant_bits,

592 + VAL_DATA_BYTES=cfg["val_data_bytes"],

593 + KEY_DATA_BYTES=key_data_bytes,

594 + META_REGION_OFFSET=meta_region_offset,

595 + NUM_SOA_FIELDS=num_soa_fields,

596 + SOA_K_NORM=soa_k_norm,

597 + SOA_V_SCALE=soa_v_scale,

598 + SOA_V_ZERO=soa_v_zero,

599 + ATTN_SCALE=scale,

600 + BLOCK_D=cfg["BLOCK_D"],

601 + BLOCK_KV=BLOCK_KV,

602 + KEY_FP8=1 if key_fp8 else 0,

603 + NORM_CORRECTION=1 if norm_correction else 0,

604 + FP8_E4B15=fp8_e4b15,

605 + USE_SINKS=1 if use_sinks else 0,

606 + num_warps=1,

607 + num_stages=1,

608 + )

609 +

610 + # Stage 2: Reduce across KV splits

611 + if output_buf is not None and output_buf.shape[0] >= B:

612 + output = output_buf[:B, :Hq, :D]

613 + else:

614 + output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)

615 + if buf_holder is not None:

616 + buf_holder._tq_output_buf = output

617 + if lse_buf is not None and lse_buf.shape[0] >= B:

618 + lse = lse_buf[:B, :Hq]

619 + else:

620 + lse = torch.empty(B, Hq, dtype=torch.float32, device=device)

621 + if buf_holder is not None:

622 + buf_holder._tq_lse_buf = lse

623 +

624 + grid2 = (B, Hq)

625 + _fwd_kernel_stage2[grid2](

626 + mid_o,

627 + output,

628 + lse,

629 + seq_lens,

630 + mid_o.stride(0),

631 + mid_o.stride(1),

632 + mid_o.stride(2),

633 + output.stride(0),

634 + output.stride(1),

635 + lse.stride(0),

636 + NUM_KV_SPLITS=NUM_KV_SPLITS,

637 + BLOCK_DV=cfg["BLOCK_D"],

638 + Lv=D,

639 + num_warps=4,

640 + num_stages=2,

641 + )

642 +

643 + return output.to(query.dtype)

vllm/v1/attention/ops/turboquant_soa/triton_turboquant_decode_v2.py new

+669

@@ -0,0 +1,669 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Optimized Triton TurboQuant decode attention (v2).

4 +

5 +FLUTE-paper optimizations applied:

6 + 1. Grouped Q heads: grid over (B, Hk, splits) instead of (B, Hq, splits).

7 + Each program loads BLOCK_M Q heads sharing a KV head into a 2D tile,

8 + enabling tl.dot on tensor cores (MFMA/WMMA) for both Q·K and P·V.

9 + 2. Vectorized pair LUT: precompute pair_table[i][j] = (T[i], T[j])

10 + offline. At runtime, extract adjacent index pairs and fetch two

11 + dequantized centroids with a single gather, halving LUT lookups.

12 + 3. exp2 instead of exp: scores pre-scaled by log2(e) so the hardware-

13 + native exp2 instruction replaces the more expensive exp.

14 + 4. Wider index extraction: for 4-bit MSE, two adjacent 4-bit indices share

15 + a byte. One byte load yields both, eliminating redundant loads.

16 + 5. Centroids pre-warmed in L1 at kernel start.

17 + 6. BLOCK_KV = TILE_SIZE raised to 16-32 (from 4), reducing loop iterations

18 + and softmax rescaling overhead.

19 +

20 +Stage 2 is reused unchanged from triton_decode_attention.py.

21 +"""

22 +

23 +import math

24 +from typing import Any

25 +

26 +import torch

27 +

28 +from vllm.platforms import current_platform

29 +from vllm.triton_utils import tl, triton

30 +from vllm.v1.attention.ops.triton_decode_attention import _fwd_kernel_stage2

31 +

32 +from .triton_turboquant_decode import _use_fp8_e4b15

33 +

34 +# ROCm prefers num_stages=1 in attention-like kernels to reduce shared-memory

35 +# pressure (mirrors the pattern used in triton_decode_attention.py and

36 +# triton_turboquant_decode.py).

37 +_is_hip = current_platform.is_rocm()

38 +

39 +# On ROCm, bf16 has the same MFMA throughput as fp16 but wider dynamic range

40 +# (8-bit exponent vs 5-bit), which is safer for attention scores. On CUDA,

41 +# fp16 tensor cores may be faster than bf16 for some shapes.

42 +_DOT_DTYPE = tl.bfloat16 if _is_hip else tl.float16

43 +

44 +

45 +# ---------------------------------------------------------------------------

46 +# Pair LUT construction (FLUTE) — called once at launcher time

47 +# ---------------------------------------------------------------------------

48 +

49 +

50 +def build_pair_lut(centroids: torch.Tensor) -> torch.Tensor:

51 + """Build vectorized pair lookup table.

52 +

53 + For N centroids, returns a [N, N, 2] float32 tensor where

54 + pair_lut[i, j] = (centroids[i], centroids[j]).

55 + Flattened to [N*N, 2] for kernel indexing: pair_lut[i*N + j].

56 +

57 + For 4-bit MSE (N=16): 16*16*2*4 = 2048 bytes — fits in L1/smem.

58 + For 3-bit MSE (N=8): 8*8*2*4 = 512 bytes.

59 + """

60 + N = centroids.shape[0]

61 + # pair_lut[i,j,0] = centroids[i], pair_lut[i,j,1] = centroids[j]

62 + c = centroids.float()

63 + lut = torch.empty(N, N, 2, dtype=torch.float32, device=centroids.device)

64 + lut[:, :, 0] = c[:, None]

65 + lut[:, :, 1] = c[None, :]

66 + return lut.reshape(N * N, 2).contiguous()

67 +

68 +

69 +# ---------------------------------------------------------------------------

70 +# Stage 1 v2: Grouped-Q + pair LUT + tl.dot + exp2

71 +# ---------------------------------------------------------------------------

72 +

73 +

74 +@triton.jit

75 +def _tq_decode_stage1_v2(

76 + Q_rot_ptr,

77 + KV_cache_ptr,

78 + KV_cache_u16_ptr, # uint16 view — Opt#3 SoA loads

79 + Block_table_ptr,

80 + Seq_lens_ptr,

81 + Centroids_ptr,

82 + Pair_lut_ptr,

83 + Mid_o_ptr,

84 + stride_qb,

85 + stride_qh,

86 + stride_cache_block, # bytes per block (bs*H*slot_aligned)

87 + stride_bt_b,

88 + stride_mid_b,

89 + stride_mid_h,

90 + stride_mid_s,

91 + NUM_Q_HEADS: tl.constexpr,

92 + NUM_KV_HEADS: tl.constexpr,

93 + HEAD_DIM: tl.constexpr,

94 + BLOCK_SIZE: tl.constexpr,

95 + PADDED_SLOT: tl.constexpr,

96 + MAX_NUM_BLOCKS: tl.constexpr,

97 + NUM_KV_SPLITS: tl.constexpr,

98 + KV_GROUP_SIZE: tl.constexpr,

99 + MSE_BITS: tl.constexpr,

100 + MSE_BYTES: tl.constexpr,

101 + VQB: tl.constexpr,

102 + VAL_DATA_BYTES: tl.constexpr,

103 + # Opt#3 SoA layout

104 + KEY_DATA_BYTES: tl.constexpr,

105 + META_REGION_OFFSET: tl.constexpr,

106 + NUM_SOA_FIELDS: tl.constexpr,

107 + SOA_K_NORM: tl.constexpr,

108 + SOA_V_SCALE: tl.constexpr,

109 + SOA_V_ZERO: tl.constexpr,

110 + N_CENTROIDS: tl.constexpr,

111 + ATTN_SCALE: tl.constexpr,

112 + BLOCK_D: tl.constexpr,

113 + TILE_SIZE: tl.constexpr,

114 + BLOCK_M: tl.constexpr,

115 + KEY_FP8: tl.constexpr,

116 + NORM_CORRECTION: tl.constexpr = 0,

117 + FP8_E4B15: tl.constexpr = 0,

118 + USE_PAIR_LUT: tl.constexpr = 0,

119 + USE_BF16_DOT: tl.constexpr = 0,

120 +):

121 + # Grid: (B, Hk, NUM_KV_SPLITS) — one program per KV head group

122 + bid = tl.program_id(0)

123 + kv_hid = tl.program_id(1)

124 + sid = tl.program_id(2)

125 +

126 + seq_len = tl.load(Seq_lens_ptr + bid)

127 + split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)

128 + split_start = split_len * sid

129 + split_end = tl.minimum(split_start + split_len, seq_len)

130 +

131 + if split_start >= split_end:

132 + return

133 +

134 + # Dimension and tile offsets

135 + d_offs = tl.arange(0, BLOCK_D)

136 + d_mask = d_offs < HEAD_DIM

137 + m_offs = tl.arange(0, BLOCK_M)

138 + t_offs = tl.arange(0, TILE_SIZE)

139 +

140 + # Q head indices for this KV group: [kv_hid * GS, kv_hid * GS + GS)

141 + q_head_offs = kv_hid * KV_GROUP_SIZE + m_offs

142 + q_mask = m_offs < KV_GROUP_SIZE

143 +

144 + # Load Q tile: [BLOCK_M, BLOCK_D] float32

145 + q_addrs = bid * stride_qb + q_head_offs[:, None] * stride_qh + d_offs[None, :]

146 + Q = tl.load(

147 + Q_rot_ptr + q_addrs,

148 + mask=q_mask[:, None] & d_mask[None, :],

149 + other=0.0,

150 + )

151 +

152 + # Pre-warm centroids in L1

153 + if not KEY_FP8:

154 + _c_warm = tl.load(Centroids_ptr + tl.arange(0, N_CENTROIDS)) # noqa: F841

155 +

156 + # Loop-invariant bit-extraction constants for MSE keys

157 + if not KEY_FP8:

158 + if MSE_BITS == 4:

159 + pass # 4-bit byte extraction is computed inline in the tile loop

160 + else:

161 + # Generic bit extraction

162 + mse_bit_off = d_offs * MSE_BITS

163 + mse_byte_idx = mse_bit_off // 8

164 + mse_bit_shift = mse_bit_off % 8

165 + mse_mask_val = (1 << MSE_BITS) - 1

166 +

167 + # Loop-invariant bit-extraction for 3-bit values

168 + if VQB == 3:

169 + val_bit_off = d_offs * 3

170 + val_byte_idx = val_bit_off // 8

171 + val_bit_shift = val_bit_off % 8

172 +

173 + # exp2 pre-scaling

174 + RCP_LN2: tl.constexpr = 1.4426950408889634

175 + LN2: tl.constexpr = 0.6931471805599453

176 + QK_SCALE = ATTN_SCALE * RCP_LN2

177 +

178 + # Online softmax accumulators: [BLOCK_M] for max, [BLOCK_M] for L

179 + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

180 + L = tl.zeros([BLOCK_M], dtype=tl.float32)

181 +

182 + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32)

183 +

184 + bt_base = bid * stride_bt_b

185 +

186 + # ================================================================

187 + # TILED LOOP: TILE_SIZE KV tokens per iteration

188 + # ================================================================

189 + for start_n in range(split_start, split_end, TILE_SIZE):

190 + kv_offs = start_n + t_offs

191 + kv_mask_1d = kv_offs < split_end

192 +

193 + page_idx = kv_offs // BLOCK_SIZE

194 + page_off = kv_offs % BLOCK_SIZE

195 + block_nums = tl.load(

196 + Block_table_ptr + bt_base + page_idx,

197 + mask=kv_mask_1d,

198 + other=0,

199 + ).to(tl.int64)

200 +

201 + # Opt#3 SoA addressing: data region then per-block SoA metadata.

202 + slot_within_block = page_off.to(tl.int64)

203 + block_base = block_nums * stride_cache_block

204 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

205 + data_bases = (

206 + block_base

207 + + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)

208 + + tl.cast(kv_hid, tl.int64) * DATA_BYTES_PER_SLOT

209 + )

210 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(

211 + kv_hid, tl.int64

212 + ) * (NUM_SOA_FIELDS * BLOCK_SIZE)

213 + knorm_u16_addrs = (

214 + head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block

215 + )

216 + vscale_u16_addrs = (

217 + head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block

218 + )

219 + vzero_u16_addrs = (

220 + head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block

221 + )

222 +

223 + # ============================================================

224 + # KEY DEQUANT → K_T: [BLOCK_D, TILE_SIZE] for tl.dot(Q, K_T)

225 + # ============================================================

226 + if KEY_FP8:

227 + k_addrs = data_bases[:, None] + d_offs[None, :]

228 + k_raw = tl.load(

229 + KV_cache_ptr + k_addrs,

230 + mask=kv_mask_1d[:, None] & d_mask[None, :],

231 + other=0,

232 + )

233 + if FP8_E4B15:

234 + K_f = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)

235 + else:

236 + K_f = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)

237 + # K_f: [TILE_SIZE, BLOCK_D] → transpose for tl.dot

238 + K_T = tl.trans(K_f)

239 + # S: [BLOCK_M, TILE_SIZE] (bf16 dot on ROCm for better dynamic range)

240 + if USE_BF16_DOT:

241 + S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))

242 + else:

243 + S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))

244 + else:

245 + # --- MSE dequantization ---

246 + if MSE_BITS == 4 and USE_PAIR_LUT:

247 + # OPT#1 (load-halving): FLUTE §3.2 pair LUT, but with each

248 + # packed byte loaded EXACTLY ONCE instead of twice (the

249 + # previous code built byte_addrs from d_offs // 2 which has

250 + # the pattern [0,0,1,1,2,2,...], duplicating every byte

251 + # load). We load the packed bytes as a [TILE_SIZE, HALF_D]

252 + # tile, decode both nibbles, then issue a single 3-D gather

253 + # against pair_lut which returns (T[lo], T[hi]) in slots

254 + # 0/1. Reshape [TILE_SIZE, HALF_D, 2] -> [TILE_SIZE, BLOCK_D]

255 + # interleaves the pairs into the expected dim order

256 + # (d=2k -> T[lo_k], d=2k+1 -> T[hi_k]).

257 + HALF_D: tl.constexpr = BLOCK_D // 2

258 + half_offs = tl.arange(0, HALF_D)

259 + # Each byte k covers dims (2k, 2k+1); it is in-range iff

260 + # 2k < HEAD_DIM. HEAD_DIM is always even in our layouts,

261 + # so this is equivalent to k < HEAD_DIM // 2.

262 + byte_mask = (half_offs * 2) < HEAD_DIM

263 +

264 + byte_addrs = data_bases[:, None] + half_offs[None, :]

265 + byte_raw = tl.load(

266 + KV_cache_ptr + byte_addrs,

267 + mask=kv_mask_1d[:, None] & byte_mask[None, :],

268 + other=0,

269 + ).to(tl.int32)

270 +

271 + lo_idx = byte_raw & 0xF # [TILE_SIZE, HALF_D]

272 + hi_idx = (byte_raw >> 4) & 0xF # [TILE_SIZE, HALF_D]

273 + pair_key = lo_idx * N_CENTROIDS + hi_idx # [TILE_SIZE, HALF_D]

274 +

275 + # Single 3-D gather: pair_lut[pair_key*2 + {0,1}]

276 + # -> [TILE_SIZE, HALF_D, 2] = (T[lo_k], T[hi_k]) per element.

277 + pair_slot = tl.arange(0, 2)

278 + c_pair = tl.load(

279 + Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],

280 + mask=(kv_mask_1d[:, None, None] & byte_mask[None, :, None]),

281 + other=0.0,

282 + )

283 + # Row-major reshape interleaves pairs:

284 + # out[t, 2k] = c_pair[t, k, 0] = T[lo_k]

285 + # out[t, 2k+1] = c_pair[t, k, 1] = T[hi_k]

286 + c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])

287 + elif MSE_BITS == 4:

288 + # 4-bit without pair LUT: exploit byte alignment

289 + half_idx = d_offs // 2

290 + nibble_shift = (d_offs % 2) * 4

291 + mse_addrs = data_bases[:, None] + half_idx[None, :]

292 + mse_raw = tl.load(

293 + KV_cache_ptr + mse_addrs,

294 + mask=kv_mask_1d[:, None] & d_mask[None, :],

295 + other=0,

296 + ).to(tl.int32)

297 + mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF

298 + c_vals = tl.load(

299 + Centroids_ptr + mse_idx,

300 + mask=kv_mask_1d[:, None] & d_mask[None, :],

301 + other=0.0,

302 + )

303 + else:

304 + # Generic 3-bit path: two byte loads → 16-bit → shift/mask

305 + mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]

306 + mse_raw0 = tl.load(

307 + KV_cache_ptr + mse_addrs0,

308 + mask=kv_mask_1d[:, None] & d_mask[None, :],

309 + other=0,

310 + ).to(tl.int32)

311 + mse_raw1 = tl.load(

312 + KV_cache_ptr + mse_addrs0 + 1,

313 + mask=kv_mask_1d[:, None] & d_mask[None, :],

314 + other=0,

315 + ).to(tl.int32)

316 + raw16 = mse_raw0 | (mse_raw1 << 8)

317 + mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val

318 + c_vals = tl.load(

319 + Centroids_ptr + mse_idx,

320 + mask=kv_mask_1d[:, None] & d_mask[None, :],

321 + other=0.0,

322 + )

323 +

324 + # Opt#1: norm-correction is pre-folded into the stored per-token

325 + # scalar at store time, so no per-tile sum+sqrt+divide here. The

326 + # stored vec_norm carries ||k||/||c_vec||; multiply by c_vals

327 + # below gives the original c_vec/||c_vec|| * ||k||.

328 +

329 + # Opt#3: K-norms from per-block SoA region. Single u16 load per

330 + # token; contiguous across tokens within a block.

331 + norm_u16 = tl.load(

332 + KV_cache_u16_ptr + knorm_u16_addrs, mask=kv_mask_1d, other=0

333 + )

334 + vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)

335 +

336 + # Reconstruct K: K_recon[t, d] = norm[t] * centroid[t, d]

337 + K_recon = c_vals * vec_norms[:, None]

338 + # K_T: [BLOCK_D, TILE_SIZE]

339 + K_T = tl.trans(K_recon)

340 +

341 + # S: [BLOCK_M, TILE_SIZE] via tensor core MMA

342 + if USE_BF16_DOT:

343 + S = QK_SCALE * tl.dot(Q.to(tl.bfloat16), K_T.to(tl.bfloat16))

344 + else:

345 + S = QK_SCALE * tl.dot(Q.to(tl.float16), K_T.to(tl.float16))

346 +

347 + # Mask out-of-range positions

348 + S = tl.where(kv_mask_1d[None, :], S, float("-inf"))

349 +

350 + # ============================================================

351 + # ONLINE SOFTMAX with exp2

352 + # ============================================================

353 + m_j = tl.maximum(M, tl.max(S, axis=1))

354 + m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

355 + P = tl.math.exp2(S - m_j[:, None])

356 + l_j = tl.sum(P, axis=1)

357 + alpha = tl.math.exp2(M - m_j)

358 +

359 + acc = acc * alpha[:, None]

360 + L = L * alpha + l_j

361 + M = m_j

362 +

363 + # ============================================================

364 + # VALUE DEQUANT → V: [TILE_SIZE, BLOCK_D] for tl.dot(P, V)

365 + # Opt#3: V data lives at data_base + KEY_DATA_BYTES; scale/zero

366 + # move to the SoA metadata region (one u16 load each per token).

367 + # ============================================================

368 + val_bases = data_bases + KEY_DATA_BYTES

369 +

370 + if VQB == 3:

371 + val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]

372 + val_raw0 = tl.load(

373 + KV_cache_ptr + val_addrs0,

374 + mask=kv_mask_1d[:, None] & d_mask[None, :],

375 + other=0,

376 + ).to(tl.int32)

377 + val_raw1 = tl.load(

378 + KV_cache_ptr + val_addrs0 + 1,

379 + mask=kv_mask_1d[:, None] & d_mask[None, :],

380 + other=0,

381 + ).to(tl.int32)

382 + raw16 = val_raw0 | (val_raw1 << 8)

383 + v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)

384 + else: # VQB == 4

385 + # OPT#3 (load-halving, value path): same pattern as OPT#1 for

386 + # keys -- load each byte exactly once as [TILE_SIZE, HALF_D],

387 + # decode both nibbles, then tl.interleave to [TILE_SIZE, BLOCK_D]

388 + # with pattern [v_lo0, v_hi0, v_lo1, v_hi1, ...].

389 + V_HALF_D: tl.constexpr = BLOCK_D // 2

390 + v_half_offs = tl.arange(0, V_HALF_D)

391 + v_byte_mask = (v_half_offs * 2) < HEAD_DIM

392 +

393 + val_addrs = val_bases[:, None] + v_half_offs[None, :]

394 + val_byte = tl.load(

395 + KV_cache_ptr + val_addrs,

396 + mask=kv_mask_1d[:, None] & v_byte_mask[None, :],

397 + other=0,

398 + ).to(tl.int32)

399 + v_lo = (val_byte & 0xF).to(tl.float32) # [T, HALF_D]

400 + v_hi = ((val_byte >> 4) & 0xF).to(tl.float32) # [T, HALF_D]

401 + v_idx = tl.interleave(v_lo, v_hi) # [T, BLOCK_D]

402 +

403 + scale_u16 = tl.load(

404 + KV_cache_u16_ptr + vscale_u16_addrs, mask=kv_mask_1d, other=0

405 + )

406 + zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=kv_mask_1d, other=0)

407 + v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)

408 + v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)

409 + V = v_idx * v_scales[:, None] + v_zeros[:, None]

410 +

411 + # P·V accumulation via tensor core MMA

412 + # P: [BLOCK_M, TILE_SIZE], V: [TILE_SIZE, BLOCK_D]

413 + if USE_BF16_DOT:

414 + acc += tl.dot(P.to(tl.bfloat16), V.to(tl.bfloat16))

415 + else:

416 + acc += tl.dot(P.to(tl.float16), V.to(tl.float16))

417 +

418 + # ================================================================

419 + # EPILOGUE: Store per-Q-head partial results for stage2

420 + # ================================================================

421 + # acc: [BLOCK_M, BLOCK_D], L: [BLOCK_M], M: [BLOCK_M]

422 + # Stage2 expects per-Q-head: mid_o[B, Hq, split, D+1]

423 + safe_L = tl.where(L > 0.0, L, 1.0)

424 + one_over_L = 1.0 / safe_L

425 + acc = acc * one_over_L[:, None]

426 +

427 + # Convert M from log2 space back to natural log for stage2 compatibility

428 + lse = M * LN2 + tl.log(safe_L)

429 +

430 + # Store each Q head's result using 2D scatter

431 + # q_head_offs: [BLOCK_M] = kv_hid * KV_GROUP_SIZE + m_offs

432 + out_addrs = (

433 + bid * stride_mid_b

434 + + q_head_offs[:, None] * stride_mid_h

435 + + sid * stride_mid_s

436 + + d_offs[None, :]

437 + )

438 + tl.store(

439 + Mid_o_ptr + out_addrs,

440 + acc,

441 + mask=q_mask[:, None] & d_mask[None, :],

442 + )

443 + # Store LSE per Q head

444 + lse_addrs = (

445 + bid * stride_mid_b + q_head_offs * stride_mid_h + sid * stride_mid_s + HEAD_DIM

446 + )

447 + tl.store(Mid_o_ptr + lse_addrs, lse, mask=q_mask)

448 +

449 +

450 +# ---------------------------------------------------------------------------

451 +# Launcher v2

452 +# ---------------------------------------------------------------------------

453 +

454 +_layout_cache: dict = {}

455 +

456 +

457 +def _get_layout(D, mse_bits, value_quant_bits, key_packed_size):

458 + key = (D, mse_bits, value_quant_bits, key_packed_size)

459 + cfg = _layout_cache.get(key)

460 + if cfg is None:

461 + val_data_bytes = math.ceil(D * value_quant_bits / 8)

462 + cfg = {

463 + "mse_bytes": math.ceil(D * mse_bits / 8),

464 + "val_data_bytes": val_data_bytes,

465 + "mse_bits": mse_bits,

466 + "n_centroids": 2**mse_bits,

467 + "BLOCK_D": triton.next_power_of_2(D),

468 + }

469 + _layout_cache[key] = cfg

470 + return cfg

471 +

472 +

473 +def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:

474 + """Return a fresh pair-LUT for ``centroids`` on each call.

475 +

476 + The LUT is tiny (N*N*2 fp32, e.g. 2KB for 4-bit MSE) so the build cost

477 + is negligible compared to attention. We avoid caching by data_ptr()

478 + because CUDA allocator memory reuse across different centroid tensors

479 + can silently return a stale LUT (subtle correctness bug). If this ever

480 + shows up on a profile, cache by a hash-of-values fingerprint instead.

481 + """

482 + return build_pair_lut(centroids)

483 +

484 +

485 +def triton_turboquant_decode_attention_v2(

486 + query: torch.Tensor,

487 + kv_cache: torch.Tensor,

488 + block_table: torch.Tensor,

489 + seq_lens: torch.Tensor,

490 + Pi: torch.Tensor,

491 + centroids: torch.Tensor,

492 + scale: float,

493 + mse_bits: int,

494 + key_packed_size: int,

495 + value_quant_bits: int,

496 + value_packed_size: int,

497 + max_seq_len: int = 0, # unused; kept for backward compatibility

498 + key_fp8: bool = False,

499 + norm_correction: bool = False,

500 + PiT: torch.Tensor | None = None,

501 + # Pre-allocated buffers (optional, required for CUDA-graph stability).

502 + mid_o_buf: torch.Tensor | None = None,

503 + output_buf: torch.Tensor | None = None,

504 + lse_buf: torch.Tensor | None = None,

505 + buf_holder: Any = None,

506 + # Fixed split count — MUST be a compile-time constant across iterations

507 + # for CUDA-graph capture/replay to work. Mirrors the v1 launcher contract.

508 + max_num_kv_splits: int = 32,

509 +) -> torch.Tensor:

510 + """Launch optimized TQ decode attention (v2 stage1 + shared stage2).

511 +

512 + Follows the same buffer-reuse + fixed-grid contract as the v1 launcher

513 + (triton_turboquant_decode_attention) so the backend can capture a CUDA

514 + graph across both versions.

515 + """

516 + B, Hq, D = query.shape

517 + Hk = kv_cache.shape[2]

518 + block_size = kv_cache.shape[1]

519 + padded_slot = kv_cache.shape[3]

520 + max_num_blocks = block_table.shape[1]

521 + n_centroids = centroids.shape[0]

522 + kv_group_size = Hq // Hk

523 + device = query.device

524 + del max_seq_len # no longer used: splits is fixed via max_num_kv_splits

525 +

526 + cfg = _get_layout(D, mse_bits, value_quant_bits, key_packed_size)

527 +

528 + # Opt#3 SoA layout constants (match store-side computation).

529 + key_data_bytes = D if key_fp8 else cfg["mse_bytes"]

530 + data_bytes_per_slot = key_data_bytes + cfg["val_data_bytes"]

531 + meta_region_offset = block_size * Hk * data_bytes_per_slot

532 + num_soa_fields = 2 if key_fp8 else 3

533 + soa_k_norm = 0

534 + soa_v_scale = 0 if key_fp8 else 1

535 + soa_v_zero = 1 if key_fp8 else 2

536 + kv_cache_u16 = kv_cache.view(torch.uint16)

537 +

538 + # Compute q_rot = q @ Pi.T

539 + if key_fp8:

540 + q_rot = query.float().contiguous()

541 + else:

542 + q_float = query.float()

543 + if PiT is None:

544 + PiT = Pi.T.contiguous()

545 + q_rot = (q_float @ PiT).contiguous()

546 +

547 + # BLOCK_M: pad KV_GROUP_SIZE to power of 2, minimum 16 for tensor cores

548 + BLOCK_M = max(16, triton.next_power_of_2(kv_group_size))

549 +

550 + # TILE_SIZE (BLOCK_KV): tokens per inner-loop iteration

551 + TILE_SIZE = 16

552 +

553 + # Fixed split count — must be constant across calls for cudagraph replay.

554 + NUM_KV_SPLITS = max_num_kv_splits

555 +

556 + # --- mid_o buffer reuse (same pattern as v1 launcher) ---

557 + if (

558 + mid_o_buf is not None

559 + and mid_o_buf.shape[0] >= B

560 + and mid_o_buf.shape[2] >= NUM_KV_SPLITS

561 + ):

562 + mid_o = mid_o_buf[:B, :Hq, :NUM_KV_SPLITS, :]

563 + else:

564 + mid_o = torch.empty(

565 + B,

566 + Hq,

567 + NUM_KV_SPLITS,

568 + D + 1,

569 + dtype=torch.float32,

570 + device=device,

571 + )

572 + if buf_holder is not None:

573 + buf_holder._tq_mid_o_buf = mid_o

574 +

575 + fp8_e4b15 = _use_fp8_e4b15(device.index or 0)

576 +

577 + # Build pair LUT for 4-bit MSE (FLUTE §3.2)

578 + use_pair_lut = mse_bits == 4 and not key_fp8

579 + pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids

580 +

581 + # Platform-dependent pipelining depth: ROCm prefers num_stages=1.

582 + stage1_num_stages = 1 if _is_hip else 2

583 +

584 + # --- Stage 1: v2 kernel ---

585 + # Grid over KV heads (not Q heads) — each program handles all Q heads

586 + grid = (B, Hk, NUM_KV_SPLITS)

587 + _tq_decode_stage1_v2[grid](

588 + q_rot,

589 + kv_cache,

590 + kv_cache_u16,

591 + block_table,

592 + seq_lens,

593 + centroids,

594 + pair_lut,

595 + mid_o,

596 + q_rot.stride(0),

597 + q_rot.stride(1),

598 + kv_cache.stride(0),

599 + block_table.stride(0),

600 + mid_o.stride(0),

601 + mid_o.stride(1),

602 + mid_o.stride(2),

603 + NUM_Q_HEADS=Hq,

604 + NUM_KV_HEADS=Hk,

605 + HEAD_DIM=D,

606 + BLOCK_SIZE=block_size,

607 + PADDED_SLOT=padded_slot,

608 + MAX_NUM_BLOCKS=max_num_blocks,

609 + NUM_KV_SPLITS=NUM_KV_SPLITS,

610 + KV_GROUP_SIZE=kv_group_size,

611 + MSE_BITS=mse_bits,

612 + MSE_BYTES=cfg["mse_bytes"],

613 + VQB=value_quant_bits,

614 + VAL_DATA_BYTES=cfg["val_data_bytes"],

615 + KEY_DATA_BYTES=key_data_bytes,

616 + META_REGION_OFFSET=meta_region_offset,

617 + NUM_SOA_FIELDS=num_soa_fields,

618 + SOA_K_NORM=soa_k_norm,

619 + SOA_V_SCALE=soa_v_scale,

620 + SOA_V_ZERO=soa_v_zero,

621 + N_CENTROIDS=n_centroids,

622 + ATTN_SCALE=scale,

623 + BLOCK_D=cfg["BLOCK_D"],

624 + TILE_SIZE=TILE_SIZE,

625 + BLOCK_M=BLOCK_M,

626 + KEY_FP8=1 if key_fp8 else 0,

627 + NORM_CORRECTION=1 if norm_correction else 0,

628 + FP8_E4B15=fp8_e4b15,

629 + USE_PAIR_LUT=1 if use_pair_lut else 0,

630 + USE_BF16_DOT=1 if _is_hip else 0,

631 + num_warps=4,

632 + num_stages=stage1_num_stages,

633 + )

634 +

635 + # --- output / lse buffer reuse (same pattern as v1 launcher) ---

636 + if output_buf is not None and output_buf.shape[0] >= B:

637 + output = output_buf[:B, :Hq, :D]

638 + else:

639 + output = torch.empty(B, Hq, D, dtype=torch.float32, device=device)

640 + if buf_holder is not None:

641 + buf_holder._tq_output_buf = output

642 + if lse_buf is not None and lse_buf.shape[0] >= B:

643 + lse = lse_buf[:B, :Hq]

644 + else:

645 + lse = torch.empty(B, Hq, dtype=torch.float32, device=device)

646 + if buf_holder is not None:

647 + buf_holder._tq_lse_buf = lse

648 +

649 + # --- Stage 2: reduce across KV splits (unchanged) ---

650 + grid2 = (B, Hq)

651 + _fwd_kernel_stage2[grid2](

652 + mid_o,

653 + output,

654 + lse,

655 + seq_lens,

656 + mid_o.stride(0),

657 + mid_o.stride(1),

658 + mid_o.stride(2),

659 + output.stride(0),

660 + output.stride(1),

661 + lse.stride(0),

662 + NUM_KV_SPLITS=NUM_KV_SPLITS,

663 + BLOCK_DV=cfg["BLOCK_D"],

664 + Lv=D,

665 + num_warps=4,

666 + num_stages=2,

667 + )

668 +

669 + return output.to(query.dtype)

vllm/v1/attention/ops/turboquant_soa/triton_turboquant_store.py new

+524

@@ -0,0 +1,524 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Fused Triton kernels for TurboQuant KV store.

4 +

5 +Two kernels:

6 +1. _tq_fused_store_fp8: FP8 key scatter + value uniform quantization.

7 +2. _tq_fused_store_mse: Fused binary-search bucketize + MSE index

8 + packing + value quantization.

9 +

10 +The launcher `triton_turboquant_store` selects the appropriate kernel.

11 +"""

12 +

13 +import math

14 +

15 +import torch

16 +

17 +from vllm.triton_utils import tl, triton

18 +

19 +from .triton_turboquant_decode import _use_fp8_e4b15

20 +

21 +# ═══════════════════════════════════════════════════════════════════════

22 +# Shared: value uniform quantization + pack + SoA scale/zero store

23 +# ═══════════════════════════════════════════════════════════════════════

24 +

25 +

26 +@triton.jit

27 +def _store_quantized_value(

28 + Value_ptr,

29 + KV_cache_ptr,

30 + KV_cache_u16_ptr, # Opt#3: u16 view for SoA scale/zero stores

31 + base, # pid * D offset into Value_ptr

32 + data_base, # byte offset into KV_cache_ptr for this slot+head's DATA region

33 + vscale_u16_addr, # u16 element index for V-scale in the SoA region

34 + vzero_u16_addr, # u16 element index for V-zero in the SoA region

35 + d_offs, # tl.arange(0, BLOCK_D)

36 + d_mask, # d_offs < D

37 + D: tl.constexpr,

38 + # Offset (bytes) from data_base to V-data region within the slot's data.

39 + # For MSE keys this equals MSE_BYTES; for FP8 keys it equals D.

40 + V_DATA_OFFSET: tl.constexpr,

41 + VQB: tl.constexpr,

42 + VAL_DATA_BYTES: tl.constexpr,

43 + BLOCK_D: tl.constexpr,

44 + BLOCK_VAL: tl.constexpr,

45 + BLOCK_GRP: tl.constexpr,

46 +):

47 + """Uniform quantization of values to VQB bits, pack, and store with

48 + scale/zero in the per-block SoA metadata region (Opt#3 layout)."""

49 + if VQB == 3:

50 + val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(

51 + tl.float32

52 + )

53 + val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)

54 + val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)

55 + v_scale = (val_max - val_min) / 7.0

56 + v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)

57 +

58 + q_vals = tl.minimum(

59 + tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7

60 + )

61 +

62 + grp_offs = tl.arange(0, BLOCK_GRP)

63 + grp_mask = grp_offs < (D // 8)

64 + q_grp = tl.reshape(q_vals, [BLOCK_GRP, 8])

65 + shifts_3bit = tl.arange(0, 8) * 3

66 + packed_24 = tl.sum(q_grp << shifts_3bit[None, :], axis=1)

67 + b0 = (packed_24 & 0xFF).to(tl.uint8)

68 + b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)

69 + b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)

70 + tl.store(

71 + KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3,

72 + b0,

73 + mask=grp_mask,

74 + )

75 + tl.store(

76 + KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3 + 1,

77 + b1,

78 + mask=grp_mask,

79 + )

80 + tl.store(

81 + KV_cache_ptr + data_base + V_DATA_OFFSET + grp_offs * 3 + 2,

82 + b2,

83 + mask=grp_mask,

84 + )

85 + else: # VQB == 4

86 + val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(

87 + tl.float32

88 + )

89 + val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)

90 + val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)

91 + v_scale = (val_max - val_min) / 15.0

92 + v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)

93 +

94 + # Quantize all D elements from register (no re-load)

95 + q_all = tl.minimum(

96 + tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 15

97 + )

98 + q_pairs = tl.reshape(q_all, [BLOCK_D // 2, 2])

99 + shifts_4 = tl.arange(0, 2) * 4

100 + packed_val = tl.sum((q_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)

101 + val_offs = tl.arange(0, BLOCK_D // 2)

102 + val_mask = val_offs < VAL_DATA_BYTES

103 + tl.store(

104 + KV_cache_ptr + data_base + V_DATA_OFFSET + val_offs,

105 + packed_val,

106 + mask=val_mask,

107 + )

108 +

109 + # ── SoA scale/zero store (Opt#3) ──────────────────────────────────

110 + # Scale and zero each get one fp16 stored at their token's slot index

111 + # within the per-head SoA region. Single u16 store per field vs the

112 + # previous 2×u8 writes per field.

113 + sc_u16 = v_scale.to(tl.float16).to(tl.uint16, bitcast=True)

114 + zr_u16 = val_min.to(tl.float16).to(tl.uint16, bitcast=True)

115 + tl.store(KV_cache_u16_ptr + vscale_u16_addr, sc_u16)

116 + tl.store(KV_cache_u16_ptr + vzero_u16_addr, zr_u16)

117 +

118 +

119 +# ═══════════════════════════════════════════════════════════════════════

120 +# FP8 key store + value uniform quantization

121 +# ═══════════════════════════════════════════════════════════════════════

122 +

123 +

124 +@triton.jit

125 +def _tq_fused_store_fp8(

126 + Key_ptr, # [NH, D] float16/bfloat16 — raw keys

127 + Value_ptr, # [NH, D] float16/bfloat16 — raw values

128 + KV_cache_ptr, # [total_bytes] uint8 (flattened view)

129 + KV_cache_u16_ptr, # uint16-aliased view of the same storage (Opt#3)

130 + Slot_mapping_ptr, # [N] int32 — per-token slot indices

131 + # Cache strides (stride_cache_block = bytes per block = bs*H*slot_aligned)

132 + stride_cache_block: tl.constexpr,

133 + # Dimensions

134 + D: tl.constexpr,

135 + H: tl.constexpr,

136 + BLOCK_SIZE: tl.constexpr,

137 + BLOCK_D: tl.constexpr,

138 + # SoA layout (Opt#3): data region is [bs, H, KEY_DATA_BYTES+VAL_DATA_BYTES],

139 + # followed by metadata region [H, NUM_SOA_FIELDS, bs] of fp16.

140 + KEY_DATA_BYTES: tl.constexpr, # FP8: = D

141 + VAL_DATA_BYTES: tl.constexpr,

142 + META_REGION_OFFSET: tl.constexpr, # bytes from block start to SoA region

143 + NUM_SOA_FIELDS: tl.constexpr, # FP8: 2 (V-scale, V-zero)

144 + SOA_V_SCALE: tl.constexpr, # FP8: 0

145 + SOA_V_ZERO: tl.constexpr, # FP8: 1

146 + # Value quantization

147 + VQB: tl.constexpr,

148 + # Packing block sizes

149 + BLOCK_VAL: tl.constexpr,

150 + BLOCK_GRP: tl.constexpr = 16,

151 + FP8_E4B15: tl.constexpr = 0, # 1 = e4b15 (Ampere/Ada), 0 = e4nv (Hopper+)

152 +):

153 + """FP8 key cast+scatter + value uniform quantization (SoA layout)."""

154 + pid = tl.program_id(0)

155 + token_idx = pid // H

156 + head_idx = pid % H

157 +

158 + slot = tl.load(Slot_mapping_ptr + token_idx)

159 + if slot < 0:

160 + return

161 + blk = (slot // BLOCK_SIZE).to(tl.int64)

162 + off = (slot % BLOCK_SIZE).to(tl.int64)

163 + head_idx_i64 = tl.cast(head_idx, tl.int64)

164 +

165 + # Block start byte address — stride_cache_block is unchanged (same total

166 + # bytes per block as the AoS layout).

167 + block_base = blk * stride_cache_block

168 +

169 + # Data region: per-slot stride = H * DATA_BYTES_PER_SLOT; per-head stride

170 + # within a slot = DATA_BYTES_PER_SLOT.

171 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

172 + data_base = (

173 + block_base

174 + + off * (H * DATA_BYTES_PER_SLOT)

175 + + head_idx_i64 * DATA_BYTES_PER_SLOT

176 + )

177 +

178 + # SoA metadata region (u16 element indexing): head strip = NUM_SOA_FIELDS

179 + # consecutive bs-long fp16 arrays.

180 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + head_idx_i64 * (

181 + NUM_SOA_FIELDS * BLOCK_SIZE

182 + )

183 + vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + off

184 + vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + off

185 +

186 + base = pid * D

187 +

188 + # ── FP8 KEY: cast to FP8 in-kernel and store at data_base ────────

189 + d_offs = tl.arange(0, BLOCK_D)

190 + d_mask = d_offs < D

191 + k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0)

192 + k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)

193 + k_bytes = k_fp8.to(tl.uint8, bitcast=True)

194 + tl.store(KV_cache_ptr + data_base + d_offs, k_bytes, mask=d_mask)

195 +

196 + # ── VALUE QUANTIZE + PACK + SoA scale/zero store ─────────────────

197 + _store_quantized_value(

198 + Value_ptr,

199 + KV_cache_ptr,

200 + KV_cache_u16_ptr,

201 + base,

202 + data_base,

203 + vscale_u16_addr,

204 + vzero_u16_addr,

205 + d_offs,

206 + d_mask,

207 + D=D,

208 + V_DATA_OFFSET=KEY_DATA_BYTES, # V follows K in the data region

209 + VQB=VQB,

210 + VAL_DATA_BYTES=VAL_DATA_BYTES,

211 + BLOCK_D=BLOCK_D,

212 + BLOCK_VAL=BLOCK_VAL,

213 + BLOCK_GRP=BLOCK_GRP,

214 + )

215 +

216 +

217 +# ═══════════════════════════════════════════════════════════════════════

218 +# Fused MSE store: bucketize + MSE index pack + norm store + value pack

219 +# (eliminates 4 PyTorch kernel launches per layer vs pack-only kernel)

220 +# ═══════════════════════════════════════════════════════════════════════

221 +

222 +

223 +@triton.jit

224 +def _tq_fused_store_mse(

225 + # Post-rotation inputs

226 + Y_ptr, # [NH, D] float32 — rotated normalized keys (x_hat @ PiT)

227 + Norms_ptr, # [NH] float32 — key vector norms (||k||)

228 + Value_ptr, # [NH, D] float32 — raw values

229 + # Quantization tables

230 + Midpoints_ptr, # [n_centroids-1] float32

231 + Centroids_ptr, # [n_centroids] float32 — used only when NORM_CORRECTION=1

232 + # Cache and indexing

233 + KV_cache_ptr, # [total_bytes] uint8 (flattened view)

234 + KV_cache_u16_ptr, # uint16-aliased view of the same storage (Opt#3)

235 + Slot_mapping_ptr, # [N] int32 — per-token slot indices

236 + # Cache strides (stride_cache_block = bs*H*slot_aligned, unchanged)

237 + stride_cache_block: tl.constexpr,

238 + # Dimensions

239 + D: tl.constexpr,

240 + H: tl.constexpr,

241 + BLOCK_SIZE: tl.constexpr,

242 + BLOCK_D: tl.constexpr,

243 + # SoA layout (Opt#3): data region [bs, H, MSE_BYTES+VAL_DATA_BYTES],

244 + # then metadata region [H, NUM_SOA_FIELDS, bs] of fp16.

245 + MSE_BYTES: tl.constexpr,

246 + VAL_DATA_BYTES: tl.constexpr,

247 + META_REGION_OFFSET: tl.constexpr, # bytes

248 + NUM_SOA_FIELDS: tl.constexpr, # MSE: 3 (k_norm, v_scale, v_zero)

249 + SOA_K_NORM: tl.constexpr, # MSE: 0

250 + SOA_V_SCALE: tl.constexpr, # MSE: 1

251 + SOA_V_ZERO: tl.constexpr, # MSE: 2

252 + # Value quantization

253 + VQB: tl.constexpr,

254 + # Packing block sizes

255 + BLOCK_VAL: tl.constexpr,

256 + # MSE params

257 + MSE_BITS: tl.constexpr,

258 + N_CENTROIDS: tl.constexpr,

259 + # When 1, fold the per-token centroid-vector norm ||c_t|| into the

260 + # stored scalar: store ||k_t|| / ||c_t|| instead of ||k_t||. This

261 + # removes the per-tile sum+sqrt+divide cost from the attention kernel's

262 + # K-load path. Load-side math becomes a plain multiply K = c * stored.

263 + NORM_CORRECTION: tl.constexpr = 0,

264 + BLOCK_GRP: tl.constexpr = 16,

265 +):

266 + """Fused MSE quantize + pack + SoA-metadata store (Opt#3 layout).

267 +

268 + Binary-search bucketize → MSE-index pack → K-norm store (SoA) → value

269 + uniform quantize + pack → V scale/zero store (SoA), in one kernel.

270 + """

271 + pid = tl.program_id(0)

272 + token_idx = pid // H

273 + head_idx = pid % H

274 +

275 + slot = tl.load(Slot_mapping_ptr + token_idx)

276 + if slot < 0:

277 + return

278 + blk = (slot // BLOCK_SIZE).to(tl.int64)

279 + off = (slot % BLOCK_SIZE).to(tl.int64)

280 + head_idx_i64 = tl.cast(head_idx, tl.int64)

281 +

282 + # Block start byte address — stride unchanged from AoS (same total

283 + # bytes per block); layout inside the block differs.

284 + block_base = blk * stride_cache_block

285 +

286 + # Data region offset (packed K indices + packed V data).

287 + DATA_BYTES_PER_SLOT: tl.constexpr = MSE_BYTES + VAL_DATA_BYTES

288 + data_base = (

289 + block_base

290 + + off * (H * DATA_BYTES_PER_SLOT)

291 + + head_idx_i64 * DATA_BYTES_PER_SLOT

292 + )

293 +

294 + # SoA metadata addresses (u16 element indexing).

295 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + head_idx_i64 * (

296 + NUM_SOA_FIELDS * BLOCK_SIZE

297 + )

298 + knorm_u16_addr = head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + off

299 + vscale_u16_addr = head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + off

300 + vzero_u16_addr = head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + off

301 +

302 + base = pid * D

303 + d_offs = tl.arange(0, BLOCK_D)

304 + d_mask = d_offs < D

305 +

306 + # ── 1. BINARY SEARCH BUCKETIZE ────────────────────────────────────

307 + y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)

308 + lo = tl.zeros([BLOCK_D], dtype=tl.int32)

309 + hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)

310 + for _ in range(MSE_BITS):

311 + mid = (lo + hi) >> 1

312 + safe_mid = tl.minimum(mid, N_CENTROIDS - 2)

313 + mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)

314 + lo = tl.where(y_vec >= mid_val, mid + 1, lo)

315 + hi = tl.where(y_vec >= mid_val, hi, mid)

316 + idx = tl.minimum(lo, N_CENTROIDS - 1)

317 +

318 + # ── 2. PACK MSE INDICES into data region at data_base[0:MSE_BYTES] ──

319 + if MSE_BITS == 4:

320 + idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])

321 + shifts_4 = tl.arange(0, 2) * 4

322 + packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)

323 + mse_offs = tl.arange(0, BLOCK_D // 2)

324 + mse_mask = mse_offs < MSE_BYTES

325 + tl.store(KV_cache_ptr + data_base + mse_offs, packed, mask=mse_mask)

326 +

327 + elif MSE_BITS == 3:

328 + grp_offs = tl.arange(0, BLOCK_GRP)

329 + grp_mask = grp_offs < (D // 8)

330 + idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])

331 + shifts_3 = tl.arange(0, 8) * 3

332 + packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)

333 + b0 = (packed_24 & 0xFF).to(tl.uint8)

334 + b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)

335 + b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)

336 + tl.store(KV_cache_ptr + data_base + grp_offs * 3, b0, mask=grp_mask)

337 + tl.store(KV_cache_ptr + data_base + grp_offs * 3 + 1, b1, mask=grp_mask)

338 + tl.store(KV_cache_ptr + data_base + grp_offs * 3 + 2, b2, mask=grp_mask)

339 +

340 + # ── 3. STORE K-norm in SoA region (single u16 store) ─────────────

341 + # Norm-correction (Opt#1) still folds 1/||c_vec|| into the stored value;

342 + # only the destination moved to the SoA K-norm array.

343 + vn_f32 = tl.load(Norms_ptr + pid) # ||k_t||, fp32

344 +

345 + if NORM_CORRECTION:

346 + safe_idx = tl.minimum(tl.maximum(idx, 0), N_CENTROIDS - 1)

347 + c_vals = tl.load(Centroids_ptr + safe_idx, mask=d_mask, other=0.0)

348 + c_norm_sq = tl.sum(tl.where(d_mask, c_vals * c_vals, 0.0), axis=0)

349 + c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)

350 + vn_f32 = vn_f32 * c_inv_norm

351 +

352 + vn_u16 = vn_f32.to(tl.float16).to(tl.uint16, bitcast=True)

353 + tl.store(KV_cache_u16_ptr + knorm_u16_addr, vn_u16)

354 +

355 + # ── 4. VALUE QUANTIZE + PACK + SoA scale/zero store ──────────────

356 + _store_quantized_value(

357 + Value_ptr,

358 + KV_cache_ptr,

359 + KV_cache_u16_ptr,

360 + base,

361 + data_base,

362 + vscale_u16_addr,

363 + vzero_u16_addr,

364 + d_offs,

365 + d_mask,

366 + D=D,

367 + V_DATA_OFFSET=MSE_BYTES, # V follows K indices in the data region

368 + VQB=VQB,

369 + VAL_DATA_BYTES=VAL_DATA_BYTES,

370 + BLOCK_D=BLOCK_D,

371 + BLOCK_VAL=BLOCK_VAL,

372 + BLOCK_GRP=BLOCK_GRP,

373 + )

374 +

375 +

376 +# ═══════════════════════════════════════════════════════════════════════

377 +# Launcher

378 +# ═══════════════════════════════════════════════════════════════════════

379 +

380 +

381 +def triton_turboquant_store(

382 + key: torch.Tensor, # [N, H, D] — raw keys (post-RoPE)

383 + value: torch.Tensor, # [N, H, D] — raw values

384 + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8

385 + slot_mapping: torch.Tensor, # [N] int32

386 + PiT: torch.Tensor, # [D, D] float32

387 + midpoints: torch.Tensor, # [n_centroids-1] float32

388 + mse_bits: int,

389 + key_packed_size: int,

390 + value_quant_bits: int,

391 + key_fp8: bool = False,

392 + centroids: torch.Tensor

393 + | None = None, # [n_centroids] float32 — required when norm_correction=True

394 + norm_correction: bool = False,

395 +):

396 + """Launch TQ store kernel (FP8 or MSE path).

397 +

398 + When ``norm_correction=True`` (default off; set by MSE norm-correction

399 + presets), the centroid-vector norm ||c_t|| is folded into the stored

400 + per-token scalar at store time (stored value = ||k_t|| / ||c_t||). The

401 + load-side attention kernel no longer computes per-tile sum+sqrt+divide,

402 + saving ~26% of TQ overhead on gpt-oss decode (ctx=4096). ``centroids``

403 + must be provided in this mode.

404 + """

405 + if norm_correction and not key_fp8:

406 + assert centroids is not None, (

407 + "norm_correction=True requires centroids tensor; got None. "

408 + "Pass the same [n_centroids] fp32 tensor used to derive midpoints."

409 + )

410 + N, H, D = key.shape

411 + NH = N * H

412 + block_size = kv_cache.shape[1]

413 + BLOCK_D = triton.next_power_of_2(D)

414 + mse_bytes = math.ceil(D * mse_bits / 8)

415 + n_centroids = 2**mse_bits

416 +

417 + val_data_bytes = math.ceil(D * value_quant_bits / 8)

418 +

419 + BLOCK_VAL = triton.next_power_of_2(val_data_bytes)

420 +

421 + # Cache strides (element_size=1 for uint8, so stride in bytes = stride()).

422 + # Only stride_cache_block (== bs*H*slot_aligned) is used under Opt#3;

423 + # per-slot and per-head strides are computed from DATA_BYTES_PER_SLOT.

424 + stride_block = kv_cache.stride(0)

425 +

426 + block_grp = triton.next_power_of_2(D // 8) if D >= 8 else 1

427 +

428 + # ── Opt#3 SoA layout constants (derived locally so the launcher API

429 + # stays unchanged). Invariant: key_packed_size + value_packed_size

430 + # == slot_size_aligned == (key_data + val_data) + (k_meta + v_meta).

431 + key_data_bytes = D if key_fp8 else mse_bytes

432 + data_bytes_per_slot = key_data_bytes + val_data_bytes

433 + # Data region occupies [0, block_size * H * data_bytes_per_slot) bytes

434 + # at the start of each block; metadata region starts at that offset.

435 + meta_region_offset = block_size * H * data_bytes_per_slot

436 + num_soa_fields = 2 if key_fp8 else 3

437 + soa_v_scale = 0 if key_fp8 else 1

438 + soa_v_zero = 1 if key_fp8 else 2

439 +

440 + # uint16-aliased view of the same storage — enables single-instruction

441 + # u16 writes for the SoA K-norm / V-scale / V-zero stores. The cache

442 + # bytes are always contiguous and 2-byte aligned under the Opt#3 layout

443 + # (MSE_BYTES, VAL_DATA_BYTES, meta_region_offset are all even).

444 + kv_cache_u16 = kv_cache.view(torch.uint16)

445 +

446 + # ── FP8 PATH: in-kernel FP8 cast + scatter via fp8 kernel ──

447 + if key_fp8:

448 + k_flat = key.reshape(NH, D).contiguous()

449 + v_flat = value.reshape(NH, D).contiguous()

450 +

451 + fp8_e4b15 = _use_fp8_e4b15(key.device.index or 0)

452 +

453 + grid = (NH,)

454 + _tq_fused_store_fp8[grid](

455 + k_flat,

456 + v_flat,

457 + kv_cache.view(-1),

458 + kv_cache_u16.view(-1),

459 + slot_mapping,

460 + stride_cache_block=stride_block,

461 + D=D,

462 + H=H,

463 + BLOCK_SIZE=block_size,

464 + BLOCK_D=BLOCK_D,

465 + KEY_DATA_BYTES=key_data_bytes,

466 + VAL_DATA_BYTES=val_data_bytes,

467 + META_REGION_OFFSET=meta_region_offset,

468 + NUM_SOA_FIELDS=num_soa_fields,

469 + SOA_V_SCALE=soa_v_scale,

470 + SOA_V_ZERO=soa_v_zero,

471 + VQB=value_quant_bits,

472 + BLOCK_VAL=BLOCK_VAL,

473 + BLOCK_GRP=block_grp,

474 + FP8_E4B15=fp8_e4b15,

475 + num_warps=4,

476 + num_stages=1,

477 + )

478 + return

479 +

480 + # ── MSE PATH: external GEMM + fused bucketize/pack kernel ──

481 + # Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)

482 + k_flat = key.float().reshape(NH, D)

483 + norms = k_flat.norm(dim=1, keepdim=True)

484 + x_hat = k_flat / (norms + 1e-8)

485 + y = x_hat @ PiT

486 +

487 + v_flat = value.float().reshape(NH, D)

488 +

489 + # When norm_correction=True, the kernel folds 1/||c_t|| into the

490 + # stored per-token K-norm so the load kernel can skip per-tile norm

491 + # math (Opt#1). centroids is only dereferenced inside the kernel when

492 + # NORM_CORRECTION=1; pass midpoints as a harmless placeholder otherwise.

493 + centroids_ptr = centroids if centroids is not None else midpoints

494 + grid = (NH,)

495 + _tq_fused_store_mse[grid](

496 + y,

497 + norms.squeeze(1),

498 + v_flat,

499 + midpoints,

500 + centroids_ptr,

501 + kv_cache.view(-1),

502 + kv_cache_u16.view(-1),

503 + slot_mapping,

504 + stride_cache_block=stride_block,

505 + D=D,

506 + H=H,

507 + BLOCK_SIZE=block_size,

508 + BLOCK_D=BLOCK_D,

509 + MSE_BYTES=mse_bytes,

510 + VAL_DATA_BYTES=val_data_bytes,

511 + META_REGION_OFFSET=meta_region_offset,

512 + NUM_SOA_FIELDS=num_soa_fields,

513 + SOA_K_NORM=0,

514 + SOA_V_SCALE=soa_v_scale,

515 + SOA_V_ZERO=soa_v_zero,

516 + VQB=value_quant_bits,

517 + BLOCK_VAL=BLOCK_VAL,

518 + MSE_BITS=mse_bits,

519 + N_CENTROIDS=n_centroids,

520 + NORM_CORRECTION=1 if norm_correction else 0,

521 + BLOCK_GRP=block_grp,

522 + num_warps=4,

523 + num_stages=1,

524 + )

vllm/v1/attention/ops/turboquant_soa/triton_turboquant_unified_attention.py new

+1350

@@ -0,0 +1,1350 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Unified (prefill + decode) Triton attention kernel for TurboQuant.

4 +

5 +Structure is ported from ``vllm/v1/attention/ops/triton_unified_attention.py``

6 +(the AITER unified attention kernel already upstreamed into vLLM). Only the K

7 +and V load sites are replaced: instead of reading raw fp16 keys/values from

8 +two contiguous caches, this kernel reads TurboQuant-packed bytes from a single

9 +combined cache and dequantizes on the fly inside the tile loop.

10 +

11 +Benefits over the v1/v2 decode-only kernels:

12 +

13 +1. GQA heads are stacked into ``BLOCK_M`` and the Q·K and P·V ops are proper

14 + ``tl.dot`` tensor-core operations (MFMA on MI300X).

15 +2. The same kernel handles decode (``BLOCK_Q=1``) and prefill

16 + (``BLOCK_Q>1``) -- no more Python per-request for-loop for continuation

17 + chunks.

18 +3. Only the subset of features exercised by the current TQ paths is kept

19 + (causal, GQA). Sinks / softcap / ALiBi / sliding-window / qq-bias /

20 + mm-prefix are deferred to follow-ups.

21 +

22 +This is an opt-in v3 path behind ``VLLM_TQ_DECODE_V3``.

23 +"""

24 +

25 +from __future__ import annotations

26 +

27 +import math

28 +from typing import Any

29 +

30 +import torch

31 +

32 +from vllm.platforms import current_platform

33 +from vllm.triton_utils import tl, triton

34 +

35 +# reduce_segments is KV-format-agnostic: by the time it runs, K/V have been

36 +# consumed and only (max, expsum, partial_output) triples remain. Reuse the

37 +# baseline's implementation verbatim to avoid code duplication.

38 +from vllm.v1.attention.ops.triton_unified_attention import reduce_segments

39 +

40 +from .triton_turboquant_decode import _use_fp8_e4b15

41 +from .triton_turboquant_decode_v2 import build_pair_lut

42 +

43 +_is_hip = current_platform.is_rocm()

44 +

45 +

46 +# ---------------------------------------------------------------------------

47 +# Helper: find which sequence in the batch a global q-block index belongs to.

48 +# Identical to the unified kernel's binary search.

49 +# ---------------------------------------------------------------------------

50 +

51 +

52 +@triton.jit

53 +def _find_seq_idx(

54 + query_start_len_ptr,

55 + target_idx,

56 + num_seqs,

57 + BLOCK_Q: tl.constexpr,

58 + use_q_block_mode: tl.constexpr,

59 +):

60 + left: tl.int32 = 0

61 + right = num_seqs

62 + while left < right:

63 + mid = (left + right) // 2

64 + val = tl.load(query_start_len_ptr + mid)

65 + mid_val = val // BLOCK_Q + mid if use_q_block_mode else val

66 + if mid_val <= target_idx:

67 + left = mid + 1

68 + else:

69 + right = mid

70 + return left - 1

71 +

72 +

73 +# ---------------------------------------------------------------------------

74 +# Fused Q-rotation prologue (Opt: eliminates the launcher-side

75 +# q_rot = (query.float() @ PiT).to(dtype).contiguous()

76 +# chain, which cost ~50-60 us/step at D=64 — all launch-bound, not compute-

77 +# bound. The fused version does one small MFMA inside the attention kernel,

78 +# once per program, before the KV-tile loop. PiT is [D,D] fp32 and fits in

79 +# L1 easily (16KB at D=64, 64KB at D=128).

80 +#

81 +# Algebra is identical to the launcher path: both compute Q_rot = Q @ PiT

82 +# with fp32 accumulate. Differences are rounding-order only (MFMA tiling

83 +# vs rocBLAS tiling), well within the 4-bit KV quant noise floor.

84 +# ---------------------------------------------------------------------------

85 +

86 +

87 +@triton.jit

88 +def _tq_fuse_q_rotation(

89 + Q, # [BLOCK_M, HEAD_SIZE_PADDED] — raw query in Q.dtype

90 + PiT_ptr,

91 + PiT_stride_0: tl.int64,

92 + PiT_stride_1: tl.int64,

93 + dim_mask, # [HEAD_SIZE_PADDED] int1 — valid head dims

94 + HEAD_SIZE_PADDED: tl.constexpr,

95 +):

96 + """Fused Q @ PiT prologue. Called once per program for the MSE-key path

97 + when the launcher has passed the raw (un-rotated) query. Returns the

98 + rotated Q in the original dtype.

99 + """

100 + d_offs = tl.arange(0, HEAD_SIZE_PADDED)

101 + pit_offsets = d_offs[:, None] * PiT_stride_0 + d_offs[None, :] * PiT_stride_1

102 + pit_mask = dim_mask[:, None] & dim_mask[None, :]

103 + PiT_tile = tl.load(PiT_ptr + pit_offsets, mask=pit_mask, other=0.0).to(tl.float32)

104 + # input_precision="ieee" pins both inputs to full fp32 MFMA (no TF32

105 + # truncation). allow_tf32 is intentionally omitted — Triton rejects

106 + # passing both. This matches the launcher's rocBLAS fp32 GEMM in

107 + # algebra; rounding-order differs, diff is <= a few fp16 ulp.

108 + Q_rot = tl.dot(

109 + Q.to(tl.float32),

110 + PiT_tile,

111 + input_precision="ieee",

112 + )

113 + return Q_rot.to(Q.dtype)

114 +

115 +

116 +# ---------------------------------------------------------------------------

117 +# TQ K-tile dequant: returns K_T : [HEAD_SIZE_PADDED, TILE_SIZE] in Q.dtype

118 +# ready for tl.dot(Q, K_T).

119 +#

120 +# Packed slot layout (per (block, pos, kv_head)) is the same as v1/v2:

121 +# MSE path: [ key_idx_bytes (MSE_BYTES) | key_norm_fp16 (2B) | ...V... ]

122 +# FP8 path: [ fp8_key_bytes (D) | ...V... ]

123 +# ---------------------------------------------------------------------------

124 +

125 +

126 +@triton.jit

127 +def _tq_load_k_tile(

128 + KV_cache_ptr,

129 + KV_cache_u16_ptr, # uint16 view of KV_cache (same storage)

130 + data_bases, # [TILE_SIZE] int64 — byte offset to each token's DATA region

131 + knorm_u16_addrs, # [TILE_SIZE] int64 — u16 element index for each token's K-norm

132 + d_offs, # [HEAD_SIZE_PADDED]

133 + d_mask, # [HEAD_SIZE_PADDED] int1

134 + tile_mask, # [TILE_SIZE] int1

135 + Centroids_ptr,

136 + Pair_lut_ptr,

137 + OUT_DTYPE: tl.constexpr, # tl.float16 or tl.bfloat16

138 + HEAD_DIM: tl.constexpr,

139 + BLOCK_D: tl.constexpr, # HEAD_SIZE_PADDED — needed for pair-LUT reshape

140 + MSE_BITS: tl.constexpr,

141 + N_CENTROIDS: tl.constexpr,

142 + KEY_FP8: tl.constexpr,

143 + USE_PAIR_LUT: tl.constexpr,

144 + NORM_CORRECTION: tl.constexpr,

145 + FP8_E4B15: tl.constexpr,

146 + TILE_SIZE: tl.constexpr,

147 +):

148 + """Load + dequantize a TILE_SIZE × HEAD_SIZE block of keys and return

149 + the transposed tile K_T : [HEAD_SIZE_PADDED, TILE_SIZE].

150 +

151 + Opt#3 SoA layout: packed K data is at `data_bases[t] + [0, MSE_BYTES)`

152 + for MSE keys (or `[0, D)` for FP8). The per-token K-norm lives in the

153 + per-block SoA metadata region; `knorm_u16_addrs` already encodes its

154 + u16 element index. For decode tiles aligned with blocks, these addresses

155 + are contiguous → one coalesced wide load replaces TILE_SIZE scattered

156 + loads (the whole point of Opt#3).

157 + """

158 + if KEY_FP8:

159 + k_addrs = data_bases[:, None] + d_offs[None, :]

160 + k_raw = tl.load(

161 + KV_cache_ptr + k_addrs,

162 + mask=tile_mask[:, None] & d_mask[None, :],

163 + other=0,

164 + )

165 + if FP8_E4B15:

166 + k_f32 = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)

167 + else:

168 + k_f32 = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)

169 + K = k_f32 # [TILE_SIZE, HEAD_SIZE_PADDED]

170 + else:

171 + # MSE path: gather packed key indices + centroid LUT.

172 + if MSE_BITS == 4 and USE_PAIR_LUT:

173 + # FLUTE pair-LUT fast path — load each packed byte once, decode

174 + # both nibbles, single 3-D gather returns (T[lo], T[hi]) per byte.

175 + HALF_D: tl.constexpr = BLOCK_D // 2

176 + half_offs = tl.arange(0, HALF_D)

177 + byte_mask = (half_offs * 2) < HEAD_DIM

178 + byte_addrs = data_bases[:, None] + half_offs[None, :]

179 + byte_raw = tl.load(

180 + KV_cache_ptr + byte_addrs,

181 + mask=tile_mask[:, None] & byte_mask[None, :],

182 + other=0,

183 + ).to(tl.int32)

184 + lo_idx = byte_raw & 0xF

185 + hi_idx = (byte_raw >> 4) & 0xF

186 + pair_key = lo_idx * N_CENTROIDS + hi_idx

187 + pair_slot = tl.arange(0, 2)

188 + c_pair = tl.load(

189 + Pair_lut_ptr + pair_key[:, :, None] * 2 + pair_slot[None, None, :],

190 + mask=(tile_mask[:, None, None] & byte_mask[None, :, None]),

191 + other=0.0,

192 + )

193 + c_vals = tl.reshape(c_pair, [TILE_SIZE, BLOCK_D])

194 + elif MSE_BITS == 4:

195 + half_idx = d_offs // 2

196 + nibble_shift = (d_offs % 2) * 4

197 + mse_addrs = data_bases[:, None] + half_idx[None, :]

198 + mse_raw = tl.load(

199 + KV_cache_ptr + mse_addrs,

200 + mask=tile_mask[:, None] & d_mask[None, :],

201 + other=0,

202 + ).to(tl.int32)

203 + mse_idx = (mse_raw >> nibble_shift[None, :]) & 0xF

204 + c_vals = tl.load(

205 + Centroids_ptr + mse_idx,

206 + mask=tile_mask[:, None] & d_mask[None, :],

207 + other=0.0,

208 + )

209 + else:

210 + # Generic bit extraction (3-bit, etc.)

211 + mse_bit_off = d_offs * MSE_BITS

212 + mse_byte_idx = mse_bit_off // 8

213 + mse_bit_shift = mse_bit_off % 8

214 + mse_mask_val = (1 << MSE_BITS) - 1

215 + mse_addrs0 = data_bases[:, None] + mse_byte_idx[None, :]

216 + mse_raw0 = tl.load(

217 + KV_cache_ptr + mse_addrs0,

218 + mask=tile_mask[:, None] & d_mask[None, :],

219 + other=0,

220 + ).to(tl.int32)

221 + mse_raw1 = tl.load(

222 + KV_cache_ptr + mse_addrs0 + 1,

223 + mask=tile_mask[:, None] & d_mask[None, :],

224 + other=0,

225 + ).to(tl.int32)

226 + raw16 = mse_raw0 | (mse_raw1 << 8)

227 + mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask_val

228 + c_vals = tl.load(

229 + Centroids_ptr + mse_idx,

230 + mask=tile_mask[:, None] & d_mask[None, :],

231 + other=0.0,

232 + )

233 +

234 + # Opt#1: 1/||c_vec|| is pre-folded into the stored K-norm at store

235 + # time, so the kernel doesn't recompute norm-correction here.

236 + # Opt#3: K-norms for a tile are contiguous in the per-block SoA

237 + # region when the tile lies within one block (always true for

238 + # aligned decode tiles with TILE_SIZE == BLOCK_SIZE) — one coalesced

239 + # u16 load replaces TILE_SIZE scattered 2-byte loads.

240 + norm_u16 = tl.load(KV_cache_u16_ptr + knorm_u16_addrs, mask=tile_mask, other=0)

241 + vec_norms = norm_u16.to(tl.float16, bitcast=True).to(tl.float32)

242 + K = c_vals * vec_norms[:, None] # [TILE_SIZE, HEAD_SIZE_PADDED]

243 +

244 + K_T = tl.trans(K.to(OUT_DTYPE)) # [HEAD_SIZE_PADDED, TILE_SIZE]

245 + _ = HEAD_DIM

246 + return K_T

247 +

248 +

249 +# ---------------------------------------------------------------------------

250 +# TQ V-tile dequant: returns V : [TILE_SIZE, HEAD_SIZE_PADDED] in Q.dtype

251 +# ready for tl.dot(P, V).

252 +# ---------------------------------------------------------------------------

253 +

254 +

255 +@triton.jit

256 +def _tq_load_v_tile(

257 + KV_cache_ptr,

258 + KV_cache_u16_ptr, # uint16 view of KV_cache (same storage)

259 + val_bases, # [TILE_SIZE] int64 — byte offset to each token's V-data

260 + vscale_u16_addrs, # [TILE_SIZE] int64 — u16 element index for V-scale

261 + vzero_u16_addrs, # [TILE_SIZE] int64 — u16 element index for V-zero

262 + d_offs,

263 + d_mask,

264 + tile_mask,

265 + OUT_DTYPE: tl.constexpr,

266 + HEAD_DIM: tl.constexpr,

267 + VQB: tl.constexpr,

268 +):

269 + """Load + dequantize a TILE_SIZE × HEAD_SIZE block of values.

270 +

271 + Opt#3 SoA layout: packed V data is at `val_bases[t] + [0, VAL_DATA_BYTES)`

272 + (V-data immediately follows K-data within the slot's data region, and

273 + `val_bases` is precomputed by the caller as `data_base + KEY_DATA_BYTES`).

274 + V-scale / V-zero live in the per-block SoA metadata region at indices

275 + `vscale_u16_addrs` / `vzero_u16_addrs`. For tiles aligned to block

276 + boundaries, those addresses are contiguous → one coalesced wide load

277 + per field instead of TILE_SIZE scattered 2-byte loads.

278 + """

279 + if VQB == 4:

280 + vb_idx = d_offs // 2

281 + vb_shift = (d_offs % 2) * 4

282 + val_addrs = val_bases[:, None] + vb_idx[None, :]

283 + val_raw = tl.load(

284 + KV_cache_ptr + val_addrs,

285 + mask=tile_mask[:, None] & d_mask[None, :],

286 + other=0,

287 + ).to(tl.int32)

288 + v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)

289 + else: # VQB == 3

290 + val_bit_off = d_offs * 3

291 + val_byte_idx = val_bit_off // 8

292 + val_bit_shift = val_bit_off % 8

293 + val_addrs0 = val_bases[:, None] + val_byte_idx[None, :]

294 + val_raw0 = tl.load(

295 + KV_cache_ptr + val_addrs0,

296 + mask=tile_mask[:, None] & d_mask[None, :],

297 + other=0,

298 + ).to(tl.int32)

299 + val_raw1 = tl.load(

300 + KV_cache_ptr + val_addrs0 + 1,

301 + mask=tile_mask[:, None] & d_mask[None, :],

302 + other=0,

303 + ).to(tl.int32)

304 + raw16 = val_raw0 | (val_raw1 << 8)

305 + v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7).to(tl.float32)

306 +

307 + # SoA scale / zero loads — coalesced on aligned decode tiles.

308 + scale_u16 = tl.load(KV_cache_u16_ptr + vscale_u16_addrs, mask=tile_mask, other=0)

309 + zero_u16 = tl.load(KV_cache_u16_ptr + vzero_u16_addrs, mask=tile_mask, other=0)

310 + v_scales = scale_u16.to(tl.float16, bitcast=True).to(tl.float32)

311 + v_zeros = zero_u16.to(tl.float16, bitcast=True).to(tl.float32)

312 +

313 + V = v_idx * v_scales[:, None] + v_zeros[:, None] # [TILE_SIZE, HEAD_SIZE_PADDED]

314 + _ = HEAD_DIM

315 + return V.to(OUT_DTYPE)

316 +

317 +

318 +# ---------------------------------------------------------------------------

319 +# Unified 2D attention kernel over a TurboQuant KV cache (causal only, GQA).

320 +# Structure follows ``kernel_unified_attention_2d`` from

321 +# triton_unified_attention.py with:

322 +# * Q-head GQA stacking into BLOCK_M

323 +# * tensor-core tl.dot(Q, K) and tl.dot(P, V)

324 +# * causal + query-padding masking

325 +# * online softmax

326 +# and K/V loads replaced by the TQ helpers above.

327 +# ---------------------------------------------------------------------------

328 +

329 +

330 +@triton.jit

331 +def kernel_tq_unified_attention_2d(

332 + output_ptr, # [num_tokens, Hq, D]

333 + query_ptr, # [num_tokens, Hq, D] — raw if FUSE_Q_ROT else rotated

334 + KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8 - TQ packed

335 + KV_cache_u16_ptr, # uint16 view of same storage — Opt#2 wide metadata loads

336 + Centroids_ptr, # [n_centroids] fp32

337 + Pair_lut_ptr, # [N*N, 2] fp32 (pair-LUT for 4-bit MSE) or dummy

338 + PiT_ptr, # [D, D] fp32 — only dereferenced when FUSE_Q_ROT == 1

339 + block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]

340 + seq_lens_ptr, # [num_seqs]

341 + query_start_len_ptr, # [num_seqs+1]

342 + sinks_ptr, # [Hq] fp32 — per-head sink logits; dereferenced only when USE_SINKS

343 + scale, # float32

344 + num_query_heads: tl.constexpr,

345 + num_queries_per_kv: tl.constexpr,

346 + block_table_stride: tl.int64,

347 + query_stride_0: tl.int64,

348 + query_stride_1: tl.int64,

349 + output_stride_0: tl.int64,

350 + output_stride_1: tl.int64,

351 + stride_cache_block: tl.int64,

352 + pit_stride_0: tl.int64,

353 + pit_stride_1: tl.int64,

354 + BLOCK_SIZE: tl.constexpr,

355 + TILE_SIZE: tl.constexpr,

356 + HEAD_SIZE: tl.constexpr,

357 + HEAD_SIZE_PADDED: tl.constexpr,

358 + BLOCK_Q: tl.constexpr,

359 + BLOCK_M: tl.constexpr,

360 + num_seqs: tl.int32,

361 + # TQ layout constants

362 + MSE_BITS: tl.constexpr,

363 + MSE_BYTES: tl.constexpr,

364 + VQB: tl.constexpr,

365 + VAL_DATA_BYTES: tl.constexpr,

366 + N_CENTROIDS: tl.constexpr,

367 + KEY_FP8: tl.constexpr,

368 + USE_PAIR_LUT: tl.constexpr,

369 + # Opt#3 SoA layout constants

370 + NUM_KV_HEADS: tl.constexpr,

371 + KEY_DATA_BYTES: tl.constexpr, # MSE_BYTES for MSE, HEAD_SIZE for FP8

372 + META_REGION_OFFSET: tl.constexpr, # bytes: bs * H * (KD+VD)

373 + NUM_SOA_FIELDS: tl.constexpr, # 3 for MSE, 2 for FP8

374 + SOA_K_NORM: tl.constexpr, # 0 for MSE, unused for FP8

375 + SOA_V_SCALE: tl.constexpr, # 1 (MSE) / 0 (FP8)

376 + SOA_V_ZERO: tl.constexpr, # 2 (MSE) / 1 (FP8)

377 + NORM_CORRECTION: tl.constexpr = 0,

378 + FP8_E4B15: tl.constexpr = 0,

379 + FUSE_Q_ROT: tl.constexpr = 0, # Fused Q @ PiT prologue when 1 (MSE path)

380 + USE_SINKS: tl.constexpr = 0, # Per-head sink logit folded into softmax denom

381 +):

382 + q_block_global_idx = tl.program_id(0)

383 + kv_head_idx = tl.program_id(1)

384 +

385 + seq_idx = _find_seq_idx(

386 + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True

387 + )

388 + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx

389 + q_block_local_idx = q_block_global_idx - q_block_start_idx

390 +

391 + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)

392 + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)

393 + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index

394 +

395 + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:

396 + return

397 +

398 + offs_m = tl.arange(0, BLOCK_M)

399 + offs_d = tl.arange(0, HEAD_SIZE_PADDED)

400 + offs_t = tl.arange(0, TILE_SIZE)

401 + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

402 +

403 + query_offset_0 = cur_batch_in_all_start_index + query_pos

404 + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv

405 + query_offset = (

406 + query_offset_0[:, None] * query_stride_0

407 + + query_offset_1[:, None] * query_stride_1

408 + + offs_d[None, :]

409 + )

410 +

411 + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)

412 + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)

413 + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)

414 +

415 + # Q : [BLOCK_M, HEAD_SIZE_PADDED]

416 + Q = tl.load(

417 + query_ptr + query_offset,

418 + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],

419 + other=0.0,

420 + )

421 +

422 + if FUSE_Q_ROT:

423 + Q = _tq_fuse_q_rotation(

424 + Q,

425 + PiT_ptr,

426 + pit_stride_0,

427 + pit_stride_1,

428 + dim_mask,

429 + HEAD_SIZE_PADDED,

430 + )

431 +

432 + block_table_offset = seq_idx * block_table_stride

433 +

434 + # Online softmax state.

435 + # When USE_SINKS=1 we fold a per-head sink logit s_h into the softmax

436 + # denominator by initializing M = s_h, L = 1.0 (== exp(s_h - s_h)).

437 + # The standard online-softmax update then correctly rescales the sink

438 + # contribution as real tokens update M, and the final acc / L gives

439 + # out_h = sum_i exp(q.k_i)*V_i / (exp(s_h) + sum_i exp(q.k_i)).

440 + # Masked rows (query_mask_1 == 0) get -inf so the sink cannot pollute

441 + # invalid (q_token, q_head) positions outside the real query range.

442 + if USE_SINKS:

443 + M = tl.load(

444 + sinks_ptr + query_offset_1, mask=query_mask_1, other=float("-inf")

445 + ).to(tl.float32)

446 + else:

447 + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

448 + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

449 + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)

450 +

451 + seq_len = tl.load(seq_lens_ptr + seq_idx)

452 + context_len = seq_len - cur_batch_query_len

453 +

454 + # Longest key prefix any query row in this q-block can attend to.

455 + max_seq_prefix_len = (

456 + context_len

457 + + q_block_local_idx * BLOCK_Q

458 + + (BLOCK_M - 1) // num_queries_per_kv

459 + + 1

460 + )

461 + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)

462 + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE)

463 +

464 + for j in range(0, num_tiles):

465 + seq_offset = j * TILE_SIZE + offs_t

466 + tile_mask = seq_offset < max_seq_prefix_len

467 +

468 + physical_block_idx = tl.load(

469 + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE

470 + ).to(tl.int64)

471 +

472 + # Opt#3 SoA addressing: compute data + SoA metadata addresses once

473 + # per tile; the load helpers below issue coalesced wide loads from

474 + # these addresses for decode tiles that align with block boundaries.

475 + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64)

476 + block_base = physical_block_idx * stride_cache_block

477 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

478 + data_bases = (

479 + block_base

480 + + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)

481 + + tl.cast(kv_head_idx, tl.int64) * DATA_BYTES_PER_SLOT

482 + )

483 + val_bases = data_bases + KEY_DATA_BYTES

484 +

485 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(

486 + kv_head_idx, tl.int64

487 + ) * (NUM_SOA_FIELDS * BLOCK_SIZE)

488 + knorm_u16_addrs = (

489 + head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block

490 + )

491 + vscale_u16_addrs = (

492 + head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block

493 + )

494 + vzero_u16_addrs = (

495 + head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block

496 + )

497 +

498 + # ---- K : [HEAD_SIZE_PADDED, TILE_SIZE] in Q.dtype ----

499 + K_T = _tq_load_k_tile(

500 + KV_cache_ptr,

501 + KV_cache_u16_ptr,

502 + data_bases,

503 + knorm_u16_addrs,

504 + offs_d,

505 + dim_mask,

506 + tile_mask,

507 + Centroids_ptr,

508 + Pair_lut_ptr,

509 + OUT_DTYPE=Q.dtype,

510 + HEAD_DIM=HEAD_SIZE,

511 + BLOCK_D=HEAD_SIZE_PADDED,

512 + MSE_BITS=MSE_BITS,

513 + N_CENTROIDS=N_CENTROIDS,

514 + KEY_FP8=KEY_FP8,

515 + USE_PAIR_LUT=USE_PAIR_LUT,

516 + NORM_CORRECTION=NORM_CORRECTION,

517 + FP8_E4B15=FP8_E4B15,

518 + TILE_SIZE=TILE_SIZE,

519 + )

520 +

521 + # ---- V : [TILE_SIZE, HEAD_SIZE_PADDED] in Q.dtype ----

522 + V = _tq_load_v_tile(

523 + KV_cache_ptr,

524 + KV_cache_u16_ptr,

525 + val_bases,

526 + vscale_u16_addrs,

527 + vzero_u16_addrs,

528 + offs_d,

529 + dim_mask,

530 + tile_mask,

531 + OUT_DTYPE=Q.dtype,

532 + HEAD_DIM=HEAD_SIZE,

533 + VQB=VQB,

534 + )

535 +

536 + # S : [BLOCK_M, TILE_SIZE]

537 + S = scale * tl.dot(Q, K_T)

538 +

539 + # Causal + query-padding mask

540 + query_abs_pos = context_len + query_pos[:, None]

541 + seq_mask = seq_offset[None, :] <= query_abs_pos

542 + S = tl.where(

543 + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,

544 + S,

545 + float("-inf"),

546 + )

547 +

548 + # Online softmax

549 + m_j = tl.maximum(M, tl.max(S, axis=1))

550 + m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

551 + P = tl.exp(S - m_j[:, None])

552 + l_j = tl.sum(P, axis=1)

553 + alpha = tl.exp(M - m_j)

554 + acc = acc * alpha[:, None]

555 + L = L * alpha + l_j

556 + M = m_j

557 +

558 + # acc += P · V

559 + acc += tl.dot(P.to(V.dtype), V)

560 +

561 + # Epilogue: normalize and store

562 + acc = acc / L[:, None]

563 +

564 + output_offset = (

565 + query_offset_0[:, None] * output_stride_0

566 + + query_offset_1[:, None] * output_stride_1

567 + + offs_d[None, :]

568 + )

569 + tl.store(

570 + output_ptr + output_offset,

571 + acc,

572 + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],

573 + )

574 +

575 +

576 +# ---------------------------------------------------------------------------

577 +# Unified 3D (split-KV) attention kernel over a TurboQuant KV cache.

578 +#

579 +# Same loop body as the 2D kernel above, but with an extra ``segm_idx`` grid

580 +# axis: each program walks only ``cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE)``

581 +# tiles of the KV history and writes a *partial* (max, expsum, acc) triple to

582 +# a scratch buffer. A subsequent ``reduce_segments`` kernel merges the

583 +# ``NUM_SEGMENTS_PER_SEQ`` partials per (q_token, q_head) via the standard

584 +# online-softmax merge.

585 +#

586 +# This gives the decode path ``num_seqs × Hk × NUM_SEGMENTS_PER_SEQ`` CTAs

587 +# instead of just ``num_seqs × Hk``, which is what v2 already uses (and what

588 +# lets v2 run flat at ~0.15 ms across KV lengths on MI300X).

589 +# ---------------------------------------------------------------------------

590 +

591 +

592 +@triton.jit

593 +def kernel_tq_unified_attention_3d(

594 + segm_output_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ, HEAD_SIZE_PADDED] fp32

595 + segm_max_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ] fp32

596 + segm_expsum_ptr, # [num_tokens, Hq, NUM_SEGMENTS_PER_SEQ] fp32

597 + query_ptr, # [num_tokens, Hq, D] — raw if FUSE_Q_ROT else rotated

598 + KV_cache_ptr, # [num_blocks, block_size, Hk, padded_slot] uint8

599 + KV_cache_u16_ptr, # uint16 view of same storage — Opt#2 wide metadata loads

600 + Centroids_ptr,

601 + Pair_lut_ptr,

602 + PiT_ptr, # [D, D] fp32 — only dereferenced when FUSE_Q_ROT == 1

603 + block_tables_ptr,

604 + seq_lens_ptr,

605 + query_start_len_ptr,

606 + sinks_ptr, # [Hq] fp32 — per-head sink logits; dereferenced only when USE_SINKS

607 + scale,

608 + num_query_heads: tl.constexpr,

609 + num_queries_per_kv: tl.constexpr,

610 + block_table_stride: tl.int64,

611 + query_stride_0: tl.int64,

612 + query_stride_1: tl.int64,

613 + stride_cache_block: tl.int64,

614 + pit_stride_0: tl.int64,

615 + pit_stride_1: tl.int64,

616 + BLOCK_SIZE: tl.constexpr,

617 + TILE_SIZE: tl.constexpr,

618 + HEAD_SIZE: tl.constexpr,

619 + HEAD_SIZE_PADDED: tl.constexpr,

620 + BLOCK_Q: tl.constexpr,

621 + BLOCK_M: tl.constexpr,

622 + num_seqs: tl.int32,

623 + NUM_SEGMENTS_PER_SEQ: tl.constexpr,

624 + # TQ layout constants (same as 2D)

625 + MSE_BITS: tl.constexpr,

626 + MSE_BYTES: tl.constexpr,

627 + VQB: tl.constexpr,

628 + VAL_DATA_BYTES: tl.constexpr,

629 + N_CENTROIDS: tl.constexpr,

630 + KEY_FP8: tl.constexpr,

631 + USE_PAIR_LUT: tl.constexpr,

632 + # Opt#3 SoA layout constants

633 + NUM_KV_HEADS: tl.constexpr,

634 + KEY_DATA_BYTES: tl.constexpr,

635 + META_REGION_OFFSET: tl.constexpr,

636 + NUM_SOA_FIELDS: tl.constexpr,

637 + SOA_K_NORM: tl.constexpr,

638 + SOA_V_SCALE: tl.constexpr,

639 + SOA_V_ZERO: tl.constexpr,

640 + NORM_CORRECTION: tl.constexpr = 0,

641 + FP8_E4B15: tl.constexpr = 0,

642 + FUSE_Q_ROT: tl.constexpr = 0, # Fused Q @ PiT prologue when 1 (MSE path)

643 + USE_SINKS: tl.constexpr = 0, # Per-head sink logit folded into softmax denom

644 +):

645 + q_block_global_idx = tl.program_id(0)

646 + kv_head_idx = tl.program_id(1)

647 + segm_idx = tl.program_id(2)

648 +

649 + seq_idx = _find_seq_idx(

650 + query_start_len_ptr, q_block_global_idx, num_seqs, BLOCK_Q, True

651 + )

652 + q_block_start_idx = tl.load(query_start_len_ptr + seq_idx) // BLOCK_Q + seq_idx

653 + q_block_local_idx = q_block_global_idx - q_block_start_idx

654 +

655 + cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)

656 + cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)

657 + cur_batch_query_len = cur_batch_in_all_stop_index - cur_batch_in_all_start_index

658 +

659 + if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:

660 + return

661 +

662 + seq_len = tl.load(seq_lens_ptr + seq_idx)

663 + tiles_per_segment = tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)

664 +

665 + # Segment empty for this sequence — skip.

666 + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:

667 + return

668 +

669 + offs_m = tl.arange(0, BLOCK_M)

670 + offs_d = tl.arange(0, HEAD_SIZE_PADDED)

671 + offs_t = tl.arange(0, TILE_SIZE)

672 + query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

673 +

674 + query_offset_0 = cur_batch_in_all_start_index + query_pos

675 + query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv

676 + query_offset = (

677 + query_offset_0[:, None] * query_stride_0

678 + + query_offset_1[:, None] * query_stride_1

679 + + offs_d[None, :]

680 + )

681 +

682 + dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)

683 + query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)

684 + query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)

685 +

686 + Q = tl.load(

687 + query_ptr + query_offset,

688 + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],

689 + other=0.0,

690 + )

691 +

692 + if FUSE_Q_ROT:

693 + Q = _tq_fuse_q_rotation(

694 + Q,

695 + PiT_ptr,

696 + pit_stride_0,

697 + pit_stride_1,

698 + dim_mask,

699 + HEAD_SIZE_PADDED,

700 + )

701 +

702 + block_table_offset = seq_idx * block_table_stride

703 +

704 + # Online softmax state. For the split-KV (3D) kernel we must fold the

705 + # sink logit into *only* segment 0's (M, L) so that the stage-2 reducer

706 + # — which is a standard online-softmax merge — counts the sink exactly

707 + # once. Segments with segm_idx > 0 keep the plain -inf / 1.0 init and

708 + # therefore carry no sink contribution; stage-2 then combines them.

709 + # This mirrors the v1 sink pattern from PR #40663 (sid==0 init trick).

710 + if USE_SINKS and segm_idx == 0:

711 + M = tl.load(

712 + sinks_ptr + query_offset_1, mask=query_mask_1, other=float("-inf")

713 + ).to(tl.float32)

714 + else:

715 + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

716 + L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

717 + acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)

718 +

719 + context_len = seq_len - cur_batch_query_len

720 + max_seq_prefix_len = (

721 + context_len

722 + + q_block_local_idx * BLOCK_Q

723 + + (BLOCK_M - 1) // num_queries_per_kv

724 + + 1

725 + )

726 + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)

727 + num_tiles = tl.cdiv(max_seq_prefix_len, TILE_SIZE)

728 +

729 + # Segment bounds: clip to the causal prefix range.

730 + tile_lo = segm_idx * tiles_per_segment

731 + tile_hi = tl.minimum((segm_idx + 1) * tiles_per_segment, num_tiles)

732 +

733 + for j in range(tile_lo, tile_hi):

734 + seq_offset = j * TILE_SIZE + offs_t

735 + tile_mask = seq_offset < max_seq_prefix_len

736 +

737 + physical_block_idx = tl.load(

738 + block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE

739 + ).to(tl.int64)

740 +

741 + # Opt#3 SoA addressing (see 2D kernel for the layout invariants).

742 + slot_within_block = (seq_offset % BLOCK_SIZE).to(tl.int64)

743 + block_base = physical_block_idx * stride_cache_block

744 + DATA_BYTES_PER_SLOT: tl.constexpr = KEY_DATA_BYTES + VAL_DATA_BYTES

745 + data_bases = (

746 + block_base

747 + + slot_within_block * (NUM_KV_HEADS * DATA_BYTES_PER_SLOT)

748 + + tl.cast(kv_head_idx, tl.int64) * DATA_BYTES_PER_SLOT

749 + )

750 + val_bases = data_bases + KEY_DATA_BYTES

751 +

752 + head_meta_u16_base = (block_base + META_REGION_OFFSET) // 2 + tl.cast(

753 + kv_head_idx, tl.int64

754 + ) * (NUM_SOA_FIELDS * BLOCK_SIZE)

755 + knorm_u16_addrs = (

756 + head_meta_u16_base + SOA_K_NORM * BLOCK_SIZE + slot_within_block

757 + )

758 + vscale_u16_addrs = (

759 + head_meta_u16_base + SOA_V_SCALE * BLOCK_SIZE + slot_within_block

760 + )

761 + vzero_u16_addrs = (

762 + head_meta_u16_base + SOA_V_ZERO * BLOCK_SIZE + slot_within_block

763 + )

764 +

765 + K_T = _tq_load_k_tile(

766 + KV_cache_ptr,

767 + KV_cache_u16_ptr,

768 + data_bases,

769 + knorm_u16_addrs,

770 + offs_d,

771 + dim_mask,

772 + tile_mask,

773 + Centroids_ptr,

774 + Pair_lut_ptr,

775 + OUT_DTYPE=Q.dtype,

776 + HEAD_DIM=HEAD_SIZE,

777 + BLOCK_D=HEAD_SIZE_PADDED,

778 + MSE_BITS=MSE_BITS,

779 + N_CENTROIDS=N_CENTROIDS,

780 + KEY_FP8=KEY_FP8,

781 + USE_PAIR_LUT=USE_PAIR_LUT,

782 + NORM_CORRECTION=NORM_CORRECTION,

783 + FP8_E4B15=FP8_E4B15,

784 + TILE_SIZE=TILE_SIZE,

785 + )

786 +

787 + V = _tq_load_v_tile(

788 + KV_cache_ptr,

789 + KV_cache_u16_ptr,

790 + val_bases,

791 + vscale_u16_addrs,

792 + vzero_u16_addrs,

793 + offs_d,

794 + dim_mask,

795 + tile_mask,

796 + OUT_DTYPE=Q.dtype,

797 + HEAD_DIM=HEAD_SIZE,

798 + VQB=VQB,

799 + )

800 +

801 + S = scale * tl.dot(Q, K_T)

802 +

803 + query_abs_pos = context_len + query_pos[:, None]

804 + seq_mask = seq_offset[None, :] <= query_abs_pos

805 + S = tl.where(

806 + query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,

807 + S,

808 + float("-inf"),

809 + )

810 +

811 + m_j = tl.maximum(M, tl.max(S, axis=1))

812 + m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

813 + P = tl.exp(S - m_j[:, None])

814 + l_j = tl.sum(P, axis=1)

815 + alpha = tl.exp(M - m_j)

816 + acc = acc * alpha[:, None]

817 + L = L * alpha + l_j

818 + M = m_j

819 +

820 + acc += tl.dot(P.to(V.dtype), V)

821 +

822 + # -------- Write partial (acc, M, L) to scratch; no normalization --------

823 + # Layout matches baseline reduce_segments so we can reuse it.

824 + segm_output_offset = (

825 + query_offset_0[:, None].to(tl.int64)

826 + * (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)

827 + + query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED)

828 + + segm_idx * HEAD_SIZE_PADDED

829 + + tl.arange(0, HEAD_SIZE_PADDED)[None, :]

830 + )

831 + tl.store(

832 + segm_output_ptr + segm_output_offset,

833 + acc,

834 + mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],

835 + )

836 + segm_offset = (

837 + query_offset_0.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)

838 + + query_offset_1 * NUM_SEGMENTS_PER_SEQ

839 + + segm_idx

840 + )

841 + tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1)

842 + tl.store(segm_expsum_ptr + segm_offset, L, mask=query_mask_0 & query_mask_1)

843 +

844 +

845 +# ---------------------------------------------------------------------------

846 +# Launcher

847 +# ---------------------------------------------------------------------------

848 +

849 +

850 +_layout_cache: dict[tuple, dict[str, int]] = {}

851 +

852 +

853 +def _get_layout(D: int, mse_bits: int, value_quant_bits: int) -> dict[str, int]:

854 + key = (D, mse_bits, value_quant_bits)

855 + cfg = _layout_cache.get(key)

856 + if cfg is None:

857 + cfg = {

858 + "mse_bytes": math.ceil(D * mse_bits / 8),

859 + "val_data_bytes": math.ceil(D * value_quant_bits / 8),

860 + "BLOCK_D": triton.next_power_of_2(D),

861 + }

862 + _layout_cache[key] = cfg

863 + return cfg

864 +

865 +

866 +def _get_pair_lut(centroids: torch.Tensor) -> torch.Tensor:

867 + """Return a fresh pair-LUT for ``centroids`` on each call.

868 +

869 + Previously cached by ``(data_ptr, device)`` — that key is unsafe

870 + because the CUDA allocator can reuse freed addresses for different

871 + centroid tensors, yielding a stale (wrong-values) LUT. The LUT is tiny

872 + (N*N*2 fp32, ~2KB at N=16) so unconditional rebuild is essentially

873 + free compared to attention work. If this shows up on a profile,

874 + replace with a hash-of-values fingerprint, not data_ptr.

875 + """

876 + return build_pair_lut(centroids)

877 +

878 +

879 +# PiT in Q's dtype cached per (PiT-id, dtype) to avoid a host cast on every

880 +# launcher call. Keyed by the fp32 PiT's data_ptr so ownership/lifetime is

881 +# tied to the caller's tensor.

882 +_pit_qdtype_cache: dict = {}

883 +

884 +

885 +def _get_pit_in_query_dtype(PiT: torch.Tensor, qdtype: torch.dtype) -> torch.Tensor:

886 + key = (PiT.data_ptr(), qdtype)

887 + cached = _pit_qdtype_cache.get(key)

888 + if cached is None:

889 + cached = PiT.to(qdtype).contiguous()

890 + _pit_qdtype_cache[key] = cached

891 + return cached

892 +

893 +

894 +def triton_turboquant_unified_attention(

895 + query: torch.Tensor, # [num_tokens, Hq, D] - fp16/bf16

896 + kv_cache: torch.Tensor, # [num_blocks, block_size, Hk, padded_slot] uint8

897 + block_table: torch.Tensor, # [num_seqs, max_num_blocks] int32

898 + seq_lens: torch.Tensor, # [num_seqs] int32

899 + query_start_loc: torch.Tensor, # [num_seqs+1] int32

900 + Pi: torch.Tensor, # [D, D] fp32

901 + centroids: torch.Tensor, # [n_centroids] fp32

902 + scale: float,

903 + mse_bits: int,

904 + key_packed_size: int,

905 + value_quant_bits: int,

906 + value_packed_size: int, # unused; kept for signature parity with v2

907 + key_fp8: bool = False,

908 + norm_correction: bool = False,

909 + PiT: torch.Tensor | None = None,

910 + output: torch.Tensor | None = None,

911 + tile_size: int | None = None,

912 + max_query_len: int | None = None,

913 + max_seq_len: int | None = None,

914 + num_kv_splits: int | None = None,

915 + force_2d: bool = False,

916 + fuse_q_rot: bool = True,

917 + sinks: torch.Tensor | None = None, # [Hq] float — per-head sink logits

918 +) -> torch.Tensor:

919 + """Launch unified TQ attention (v3).

920 +

921 + ``query`` carries *raw* query vectors. For the MSE-key path the query

922 + has to be rotated by ``PiT`` before it can be multiplied against the

923 + (already-rotated) stored K. By default (``fuse_q_rot=True``) that

924 + rotation is done inside the attention kernel prologue as a single

925 + small MFMA — no extra dispatch, no HBM round-trip. Setting

926 + ``fuse_q_rot=False`` restores the original launcher path (fp32 rocBLAS

927 + GEMM + casts + .contiguous()), which is kept as an A/B toggle for

928 + bench harnesses; numerical results match the fused path to within a

929 + few ulp of fp32 round-off. The FP8-key path never rotates Q regardless

930 + of this flag.

931 +

932 + ``tile_size`` defaults to ``32`` for prefill (``max_query_len > 1``) and

933 + ``16`` for pure decode (``max_query_len == 1``). Callers may override.

934 +

935 + Dispatch rule:

936 + * Prefill / chunked-prefill (any query block with ``BLOCK_Q > 1``) and

937 + ``num_tokens > num_seqs``: 2D kernel (plenty of CTAs already).

938 + * Pure decode with long KV: 3D split-KV kernel + ``reduce_segments``

939 + for extra parallelism; this is the same pattern v2 uses.

940 + * ``force_2d=True``: force the 2D path regardless (used by the

941 + apples-to-apples bench to isolate pure TQ dequant overhead).

942 +

943 + ``num_kv_splits`` controls the 3D-split segment count (default ``16``).

944 +

945 + Returns ``output`` of shape ``[num_tokens, Hq, D]`` in ``query.dtype``.

946 + """

947 + assert query.dim() == 3, f"query must be [N, Hq, D], got {query.shape}"

948 + num_tokens, Hq, D = query.shape

949 + Hk = kv_cache.shape[2]

950 + block_size = kv_cache.shape[1]

951 + kv_group_size = Hq // Hk

952 + num_seqs = int(query_start_loc.shape[0] - 1)

953 + device = query.device

954 +

955 + cfg = _get_layout(D, mse_bits, value_quant_bits)

956 + _ = value_packed_size # unused

957 +

958 + # Q-rotation strategy:

959 + # * FP8-key path: no rotation at all (keys stored as fp8, no codebook).

960 + # * MSE-key path + fuse_q_rot=True (default): pass the raw query to the

961 + # kernel together with PiT and let the kernel's prologue do one

962 + # fp32 tl.dot(Q, PiT). Eliminates the launcher-side

963 + # (cast, rocBLAS GEMM, cast, contiguous) chain (~50-60 us/step at

964 + # D=64, all launch-bound — see bottleneck report §10).

965 + # * MSE-key path + fuse_q_rot=False: legacy launcher rotation. Kept

966 + # as an A/B toggle for bench harnesses.

967 + #

968 + # PiT is always materialized as fp32 D x D contiguous; the kernel loads

969 + # it as fp32 (cheap, 16KB at D=64) and computes Q @ PiT with IEEE fp32

970 + # accumulate via tl.dot.

971 + if key_fp8:

972 + q_rot = query.contiguous()

973 + apply_fuse_q_rot = False

974 + else:

975 + if PiT is None:

976 + PiT = Pi.T.contiguous()

977 + apply_fuse_q_rot = bool(fuse_q_rot)

978 + if apply_fuse_q_rot:

979 + q_rot = query.contiguous()

980 + else:

981 + q_rot = (query.float() @ PiT).to(query.dtype).contiguous()

982 +

983 + # PiT in fp32, contiguous. For the fused path this is what the kernel

984 + # loads; for the legacy path it's passed through as a harmless tensor

985 + # (never dereferenced under the FUSE_Q_ROT constexpr guard). On FP8

986 + # path Pi may be unused at the call site, so we fall back to an

987 + # arbitrary non-null tensor (reuse centroids) to satisfy Triton's

988 + # non-null pointer requirement.

989 + if (not key_fp8) and PiT is not None:

990 + PiT_f32 = PiT if PiT.dtype == torch.float32 else PiT.to(torch.float32)

991 + if not PiT_f32.is_contiguous():

992 + PiT_f32 = PiT_f32.contiguous()

993 + pit_stride_0 = PiT_f32.stride(0)

994 + pit_stride_1 = PiT_f32.stride(1)

995 + else:

996 + PiT_f32 = centroids # harmless dummy; not dereferenced when FUSE_Q_ROT=0

997 + pit_stride_0 = 0

998 + pit_stride_1 = 0

999 +

1000 + # Sinks: per-head fp32 logits, contiguous. The kernel only dereferences

1001 + # this pointer when USE_SINKS=1, so when the caller passes None we bind

1002 + # a harmless non-null tensor (centroids) to satisfy Triton's non-null

1003 + # pointer requirement. See the sink design notes in the 2D/3D kernel

1004 + # bodies and the v1 precedent from PR #40663 (sid==0 init trick).

1005 + if sinks is not None:

1006 + sinks_f32 = sinks if sinks.dtype == torch.float32 else sinks.to(torch.float32)

1007 + if not sinks_f32.is_contiguous():

1008 + sinks_f32 = sinks_f32.contiguous()

1009 + assert sinks_f32.numel() == Hq, (

1010 + f"sinks must have shape [Hq={Hq}], got numel={sinks_f32.numel()}"

1011 + )

1012 + use_sinks = True

1013 + else:

1014 + sinks_f32 = centroids # harmless dummy; not dereferenced when USE_SINKS=0

1015 + use_sinks = False

1016 +

1017 + if output is None:

1018 + output = torch.empty_like(query)

1019 +

1020 + # BLOCK_M heuristic (TQ-specific; diverges from stock unified_attention).

1021 + #

1022 + # Stock fp16 picks BLOCK_M=16 for small GQA because fp16 is ALU-saturated

1023 + # at any tile size and smaller BLOCK_M = better occupancy. TQ is a

1024 + # DIFFERENT story: the dequant chain (packed-byte loads + centroid gather

1025 + # + pair-LUT + norm/scale bitcasts) has big per-block fixed overhead, and

1026 + # the pair-LUT gather pattern is latency-bound (NOT HBM-bandwidth-bound,

1027 + # confirmed via rocprof: MemUnitStalled=0.01%, VALUBusy=56% at BM=16).

1028 + #

1029 + # For prefill/chunked (max_query_len > 1) we want BLOCK_M as large as

1030 + # possible to (a) amortize the fixed dequant overhead over more query

1031 + # tokens per CTA and (b) let MFMA saturate with a 128x32 tile shape.

1032 + # Empirical sweep on MI300X across gpt-oss (D=64) and llama (D=128):

1033 + # prefill B=1 Q=4k : BM=128 is 3.4x faster than BM=16

1034 + # chunked B=64 Q=1k C=8: BM=128 is 5.6x faster than BM=16

1035 + # With BM=128 v3 matches or BEATS the fp16 baseline on prefill/chunked.

1036 + #

1037 + # For pure decode (max_query_len == 1) BLOCK_M stays small because

1038 + # BLOCK_Q > 1 just pads rows — each sequence contributes at most 1 query

1039 + # token and larger BLOCK_M wastes lanes on masked-out rows.

1040 + if max_query_len is not None:

1041 + is_prefill_like = max_query_len > 1

1042 + else:

1043 + is_prefill_like = num_tokens > num_seqs

1044 +

1045 + if is_prefill_like:

1046 + # Prefill / chunked: use BLOCK_M=128 (rounded up to a multiple of

1047 + # kv_group_size). Each CTA does very large work (BLOCK_M x TILE_SIZE

1048 + # x KV_len MMA + dequant), so even small grids (~17 CTAs) saturate

1049 + # MI300X just fine. We verified empirically: BM=128 wins across all

1050 + # prefill/chunked sweep points, including small (B=1 Q=256 C=8k) and

1051 + # large (B=64 Q=1k C=8k; 5.6x speedup vs BM=16). BLOCK_M must be a

1052 + # multiple of kv_group_size so BLOCK_Q is an integer.

1053 + BLOCK_M = max(128, triton.next_power_of_2(kv_group_size))

1054 + else:

1055 + # Decode: BLOCK_Q > 1 in decode just pads rows (at most 1 query

1056 + # token per sequence), so there's little to amortize. Inherit

1057 + # stock's small-BLOCK_M heuristic.

1058 + BLOCK_M = 16 if kv_group_size <= 16 else triton.next_power_of_2(kv_group_size)

1059 + BLOCK_Q = BLOCK_M // kv_group_size

1060 +

1061 + # Grid: at most ceil(N / BLOCK_Q) + num_seqs q-blocks total, like unified.

1062 + total_num_q_blocks = num_tokens // BLOCK_Q + num_seqs

1063 +

1064 + # TILE_SIZE heuristic (matches stock unified_attention's _get_tile_size):

1065 + # prefill (max_query_len > 1): 32

1066 + # decode (max_query_len == 1): 16

1067 + if tile_size is None:

1068 + tile_size = 32 if is_prefill_like else 16

1069 +

1070 + # Pair-LUT fast path for 4-bit MSE keys. Skipped for FP8 and non-4-bit.

1071 + # When USE_PAIR_LUT==0 the kernel never dereferences pair_lut, but Triton

1072 + # still requires a tensor pointer, so we fall back to reusing `centroids`.

1073 + use_pair_lut = (not key_fp8) and (mse_bits == 4)

1074 + pair_lut = _get_pair_lut(centroids) if use_pair_lut else centroids

1075 +

1076 + fp8_e4b15 = _use_fp8_e4b15(device.index or 0)

1077 + num_stages = 1 if _is_hip else 2

1078 +

1079 + # uint16-aliased view of the cache. Under the Opt#3 SoA layout, K-norm /

1080 + # V-scale / V-zero live in a contiguous per-block metadata region and are

1081 + # always 2-byte aligned, so single-instruction u16 loads replace the

1082 + # original 2× uint8 + OR sequence for every per-token metadata fetch.

1083 + kv_cache_u16 = kv_cache.view(torch.uint16)

1084 +

1085 + # Opt#3 SoA layout constants (derived locally; matches the store-side

1086 + # computation so the launcher signature stays unchanged). Invariant:

1087 + # data_bytes_per_slot + meta_bytes_per_slot == slot_size_aligned.

1088 + mse_bytes = cfg["mse_bytes"]

1089 + val_data_bytes = cfg["val_data_bytes"]

1090 + key_data_bytes = D if key_fp8 else mse_bytes

1091 + data_bytes_per_slot = key_data_bytes + val_data_bytes

1092 + meta_region_offset = block_size * Hk * data_bytes_per_slot

1093 + num_soa_fields = 2 if key_fp8 else 3

1094 + soa_k_norm = 0 # unused for FP8; harmless constant

1095 + soa_v_scale = 0 if key_fp8 else 1

1096 + soa_v_zero = 1 if key_fp8 else 2

1097 +

1098 + # ------------------------------------------------------------------

1099 + # Dispatch: 2D for prefill / chunked; 3D split-KV for pure decode with

1100 + # long KV. Falls back to 2D if the sequences are short or if the caller

1101 + # forces it. ``is_prefill_like`` was already computed above for the

1102 + # BLOCK_M / tile_size heuristic; reuse it here.

1103 + # ------------------------------------------------------------------

1104 +

1105 + # Use 3D only for pure decode where the KV history is long enough that

1106 + # splitting meaningfully increases SM occupancy. Threshold is empirically

1107 + # tuned on MI300X:

1108 + # seq < 1024 -> 2D path (short KV, reduce overhead > parallelism gain)

1109 + # seq >= 1024 -> 3D path (v3/v2 ratio 1.04-1.27x across shapes)

1110 + # Below 1024 KV tokens the reduce_segments + scratch-alloc overhead

1111 + # exceeds the parallelism gain from splitting. Above it, the single-CTA

1112 + # 2D path is serial over the KV history and starves MI300X's 304 CUs.

1113 + #

1114 + # IMPORTANT: max_seq_len is expected to come from the caller (backends

1115 + # track it from the block table). We avoid calling seq_lens.max().item()

1116 + # here because that forces a GPU->CPU sync on every launcher call, which

1117 + # adds ~150us of overhead to pure-decode steps.

1118 + if max_seq_len is None:

1119 + # Fallback: infer from block_table (host-side, no sync). This is an

1120 + # upper bound, sufficient for the dispatch decision.

1121 + max_seq_len_hint = int(block_table.shape[1]) * int(block_size)

1122 + else:

1123 + max_seq_len_hint = int(max_seq_len)

1124 + use_3d = (not force_2d) and (not is_prefill_like) and max_seq_len_hint >= 1024

1125 +

1126 + if not use_3d:

1127 + kernel_tq_unified_attention_2d[(total_num_q_blocks, Hk)](

1128 + output_ptr=output,

1129 + query_ptr=q_rot,

1130 + KV_cache_ptr=kv_cache,

1131 + KV_cache_u16_ptr=kv_cache_u16,

1132 + Centroids_ptr=centroids,

1133 + Pair_lut_ptr=pair_lut,

1134 + PiT_ptr=PiT_f32,

1135 + block_tables_ptr=block_table,

1136 + seq_lens_ptr=seq_lens,

1137 + query_start_len_ptr=query_start_loc,

1138 + sinks_ptr=sinks_f32,

1139 + scale=scale,

1140 + num_query_heads=Hq,

1141 + num_queries_per_kv=kv_group_size,

1142 + block_table_stride=block_table.stride(0),

1143 + query_stride_0=q_rot.stride(0),

1144 + query_stride_1=q_rot.stride(1),

1145 + output_stride_0=output.stride(0),

1146 + output_stride_1=output.stride(1),

1147 + stride_cache_block=kv_cache.stride(0),

1148 + pit_stride_0=pit_stride_0,

1149 + pit_stride_1=pit_stride_1,

1150 + BLOCK_SIZE=block_size,

1151 + TILE_SIZE=tile_size,

1152 + HEAD_SIZE=D,

1153 + HEAD_SIZE_PADDED=cfg["BLOCK_D"],

1154 + BLOCK_Q=BLOCK_Q,

1155 + BLOCK_M=BLOCK_M,

1156 + num_seqs=num_seqs,

1157 + MSE_BITS=mse_bits,

1158 + MSE_BYTES=mse_bytes,

1159 + VQB=value_quant_bits,

1160 + VAL_DATA_BYTES=val_data_bytes,

1161 + N_CENTROIDS=int(centroids.numel()),

1162 + KEY_FP8=1 if key_fp8 else 0,

1163 + USE_PAIR_LUT=1 if use_pair_lut else 0,

1164 + NUM_KV_HEADS=Hk,

1165 + KEY_DATA_BYTES=key_data_bytes,

1166 + META_REGION_OFFSET=meta_region_offset,

1167 + NUM_SOA_FIELDS=num_soa_fields,

1168 + SOA_K_NORM=soa_k_norm,

1169 + SOA_V_SCALE=soa_v_scale,

1170 + SOA_V_ZERO=soa_v_zero,

1171 + NORM_CORRECTION=1 if norm_correction else 0,

1172 + FP8_E4B15=fp8_e4b15,

1173 + FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,

1174 + USE_SINKS=1 if use_sinks else 0,

1175 + num_warps=4,

1176 + num_stages=num_stages,

1177 + )

1178 + return output

1179 +

1180 + # ---------------- 3D split-KV path ----------------

1181 + # Pick segment count. 16 is a good default for MI300X (304 CUs): at

1182 + # Hk=8, num_seqs=1, total CTAs = 1 * 8 * 16 = 128, which saturates when

1183 + # combined with warp-level parallelism. Capped at ceil(max_seq_len /

1184 + # TILE_SIZE) so we don't launch empty segments.

1185 + if num_kv_splits is None:

1186 + num_kv_splits = 16

1187 + max_possible_splits = max(1, (max_seq_len_hint + tile_size - 1) // tile_size)

1188 + num_segments = max(1, min(num_kv_splits, max_possible_splits))

1189 +

1190 + HEAD_SIZE_PADDED = cfg["BLOCK_D"]

1191 + # Scratch buffers for segment partials. fp32 to match the kernel's

1192 + # accumulator dtype. Allocated per-call for now; a production backend

1193 + # would hoist these to persistent buffers reused across layers.

1194 + segm_output = torch.empty(

1195 + (num_tokens, Hq, num_segments, HEAD_SIZE_PADDED),

1196 + dtype=torch.float32,

1197 + device=device,

1198 + )

1199 + segm_max = torch.empty(

1200 + (num_tokens, Hq, num_segments), dtype=torch.float32, device=device

1201 + )

1202 + segm_expsum = torch.empty(

1203 + (num_tokens, Hq, num_segments), dtype=torch.float32, device=device

1204 + )

1205 +

1206 + kernel_tq_unified_attention_3d[(total_num_q_blocks, Hk, num_segments)](

1207 + segm_output_ptr=segm_output,

1208 + segm_max_ptr=segm_max,

1209 + segm_expsum_ptr=segm_expsum,

1210 + query_ptr=q_rot,

1211 + KV_cache_ptr=kv_cache,

1212 + KV_cache_u16_ptr=kv_cache_u16,

1213 + Centroids_ptr=centroids,

1214 + Pair_lut_ptr=pair_lut,

1215 + PiT_ptr=PiT_f32,

1216 + block_tables_ptr=block_table,

1217 + seq_lens_ptr=seq_lens,

1218 + query_start_len_ptr=query_start_loc,

1219 + sinks_ptr=sinks_f32,

1220 + scale=scale,

1221 + num_query_heads=Hq,

1222 + num_queries_per_kv=kv_group_size,

1223 + block_table_stride=block_table.stride(0),

1224 + query_stride_0=q_rot.stride(0),

1225 + query_stride_1=q_rot.stride(1),

1226 + stride_cache_block=kv_cache.stride(0),

1227 + pit_stride_0=pit_stride_0,

1228 + pit_stride_1=pit_stride_1,

1229 + BLOCK_SIZE=block_size,

1230 + TILE_SIZE=tile_size,

1231 + HEAD_SIZE=D,

1232 + HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,

1233 + BLOCK_Q=BLOCK_Q,

1234 + BLOCK_M=BLOCK_M,

1235 + num_seqs=num_seqs,

1236 + NUM_SEGMENTS_PER_SEQ=num_segments,

1237 + MSE_BITS=mse_bits,

1238 + MSE_BYTES=mse_bytes,

1239 + VQB=value_quant_bits,

1240 + VAL_DATA_BYTES=val_data_bytes,

1241 + N_CENTROIDS=int(centroids.numel()),

1242 + KEY_FP8=1 if key_fp8 else 0,

1243 + USE_PAIR_LUT=1 if use_pair_lut else 0,

1244 + NUM_KV_HEADS=Hk,

1245 + KEY_DATA_BYTES=key_data_bytes,

1246 + META_REGION_OFFSET=meta_region_offset,

1247 + NUM_SOA_FIELDS=num_soa_fields,

1248 + SOA_K_NORM=soa_k_norm,

1249 + SOA_V_SCALE=soa_v_scale,

1250 + SOA_V_ZERO=soa_v_zero,

1251 + NORM_CORRECTION=1 if norm_correction else 0,

1252 + FP8_E4B15=fp8_e4b15,

1253 + FUSE_Q_ROT=1 if apply_fuse_q_rot else 0,

1254 + USE_SINKS=1 if use_sinks else 0,

1255 + num_warps=4,

1256 + num_stages=num_stages,

1257 + )

1258 +

1259 + # Reduce: merge num_segments partials per (q_token, q_head) into the

1260 + # final output via online-softmax merge. Reusing the baseline kernel

1261 + # (KV-format-agnostic).

1262 + reduce_segments[(num_tokens, Hq)](

1263 + output_ptr=output,

1264 + segm_output_ptr=segm_output,

1265 + segm_max_ptr=segm_max,

1266 + segm_expsum_ptr=segm_expsum,

1267 + seq_lens_ptr=seq_lens,

1268 + num_seqs=num_seqs,

1269 + num_query_heads=Hq,

1270 + out_scale_inv=1.0,

1271 + output_stride_0=output.stride(0),

1272 + output_stride_1=output.stride(1),

1273 + block_table_stride=block_table.stride(0),

1274 + TILE_SIZE=tile_size,

1275 + HEAD_SIZE=D,

1276 + HEAD_SIZE_PADDED=HEAD_SIZE_PADDED,

1277 + query_start_len_ptr=query_start_loc,

1278 + BLOCK_Q=BLOCK_Q,

1279 + NUM_SEGMENTS_PER_SEQ=num_segments,

1280 + USE_FP8=False,

1281 + )

1282 + return output

1283 +

1284 +

1285 +# ---------------------------------------------------------------------------

1286 +# Thin decode adapter: lets callers who use the AoS decode launcher signature

1287 +# switch to the SoA decode without rebuilding query_start_loc themselves.

1288 +# ---------------------------------------------------------------------------

1289 +

1290 +

1291 +def triton_turboquant_decode_attention_soa(

1292 + query: torch.Tensor, # [B, Hq, D]

1293 + kv_cache: torch.Tensor,

1294 + block_table: torch.Tensor,

1295 + seq_lens: torch.Tensor,

1296 + Pi: torch.Tensor,

1297 + centroids: torch.Tensor,

1298 + scale: float,

1299 + mse_bits: int,

1300 + key_packed_size: int,

1301 + value_quant_bits: int,

1302 + value_packed_size: int,

1303 + key_fp8: bool = False,

1304 + norm_correction: bool = False,

1305 + PiT: torch.Tensor | None = None,

1306 + # kept for signature parity with v1/v2 decode launchers; unused here.

1307 + max_seq_len: int = 0,

1308 + mid_o_buf: Any = None,

1309 + output_buf: torch.Tensor | None = None,

1310 + lse_buf: Any = None,

1311 + buf_holder: Any = None,

1312 + max_num_kv_splits: int = 32,

1313 + sinks: torch.Tensor | None = None,

1314 +) -> torch.Tensor:

1315 + """Decode-only convenience wrapper around ``triton_turboquant_unified_attention``.

1316 +

1317 + Treats a rank-3 ``[B, Hq, D]`` query as one token per request (query_len = 1)

1318 + and synthesizes a ``query_start_loc`` of ``[0, 1, 2, ..., B]``.

1319 + ``max_num_kv_splits`` is forwarded to the unified launcher as the 3D

1320 + split-KV segment count (capped per-call against ``max_seq_len``).

1321 + ``sinks`` (optional ``[Hq]`` fp32) are forwarded to the kernel which

1322 + folds them into the softmax denominator via the init-time trick.

1323 + """

1324 + del mid_o_buf, lse_buf, buf_holder

1325 + B = query.shape[0]

1326 + cu_seqlens_q = torch.arange(B + 1, device=query.device, dtype=seq_lens.dtype)

1327 +

1328 + out = triton_turboquant_unified_attention(

1329 + query=query.contiguous(),

1330 + kv_cache=kv_cache,

1331 + block_table=block_table,

1332 + seq_lens=seq_lens,

1333 + query_start_loc=cu_seqlens_q,

1334 + Pi=Pi,

1335 + centroids=centroids,

1336 + scale=scale,

1337 + mse_bits=mse_bits,

1338 + key_packed_size=key_packed_size,

1339 + value_quant_bits=value_quant_bits,

1340 + value_packed_size=value_packed_size,

1341 + key_fp8=key_fp8,

1342 + norm_correction=norm_correction,

1343 + PiT=PiT,

1344 + output=output_buf[:B] if output_buf is not None else None,

1345 + max_query_len=1,

1346 + max_seq_len=max_seq_len if max_seq_len > 0 else None,

1347 + num_kv_splits=max_num_kv_splits,

1348 + sinks=sinks,

1349 + )

1350 + return out