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