| 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 |