From e895c9c008428ad061d19f7d91d42b78659842d4 Mon Sep 17 00:00:00 2001 From: Ling Zhang <69022634+ZLkanyo009@users.noreply.github.com> Date: Wed, 28 Oct 2026 08:32:24 +0000 Subject: [PATCH 1/2] [FlyDSL] pa_decode: zero V rows past the context before PV Pages past a sequence are pinned to block 0 and the unwritten tail of its last page is read as is. Their probabilities are exactly zero, but 0 * NaN is NaN in the PV MFMA, so a NaN in any such row (a recycled block 0, a checkpoint image, a freed page) turned the whole output row into NaN. Per-token V scales had the same hole through P * scale. Only a sequence's last 256-token tile holds such rows. On that tile each lane builds 64-bit byte masks from the context length once, and every head chunk ANDs its V operands with them right before its PV, behind a uniform branch, so full tiles run unchanged and the loads stay pipelined. Per-token K/V scales at or past the context are staged as zero. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/pa_decode/op_lds.py | 108 +++++++++++++-- .../ops/flydsl/kernels/pa_decode/pipeline.py | 28 +++- aiter/ops/flydsl/kernels/pa_decode_kernel.py | 9 +- op_tests/test_flydsl_pa_decode.py | 125 ++++++++++++++++++ 4 files changed, 253 insertions(+), 17 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py b/aiter/ops/flydsl/kernels/pa_decode/op_lds.py index b29bb626f3f..6b326acf45a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py +++ b/aiter/ops/flydsl/kernels/pa_decode/op_lds.py @@ -498,7 +498,23 @@ def scale_buffer_offset(self, tt_val): return (tt_val & fx.Int32(1)) * self.traits.KV_BUF_STRIDE return 0 - def stage_scales(self, phys_vec, buf_off=0): + def _context_scales(self, scales, remaining): + """Zero scales of tokens at or past the context length. + + P is zero there, but a NaN or infinite V scale still makes P * scale + NaN. ``remaining`` counts valid tokens from the vector's first element. + """ + zero = fx.Float32(0.0) + return fx.Vector.from_elements( + [ + (remaining > i).select(fx.Float32(scales[i]), zero) + for i in range_constexpr(scales.shape[0]) + ], + dtype=fx.Float32, + ) + + def stage_scales(self, phys_vec, buf_off, tt): + tile_valid = self.ctx.context_len - tt * self.traits.TILE_TOK if const_expr(self.traits.block_size >= 64): phys = fx.Int32(phys_vec[0]) chunk_tok = ( @@ -514,11 +530,11 @@ def stage_scales(self, phys_vec, buf_off=0): + self.ctx.kv_h * self.ctx.stride_ks_head + page_tok ) - k_scale_vec = self.load_k_scale(scale_idx) - v_scale_vec = self.load_v_scale(scale_idx) - slot = ( - self.ctx.warp * self.traits.TOK_PER_WARP + chunk_tok - ) * self.traits.f32 + tile_tok = self.ctx.warp * self.traits.TOK_PER_WARP + chunk_tok + remaining = tile_valid - tile_tok + k_scale_vec = self._context_scales(self.load_k_scale(scale_idx), remaining) + v_scale_vec = self._context_scales(self.load_v_scale(scale_idx), remaining) + slot = tile_tok * self.traits.f32 self.lds.store( self.traits.sKScale_off + buf_off + slot, fx.Float32, k_scale_vec ) @@ -536,11 +552,15 @@ def stage_scales(self, phys_vec, buf_off=0): k_scale_scalar = fx.Float32(self.load_k_scale(scale_idx)[0]) v_scale_scalar = fx.Float32(self.load_v_scale(scale_idx)[0]) fx.rocdl.sched_barrier(fx.rocdl.mask_vmem_rd) - slot = ( + tile_tok = ( self.ctx.warp * self.traits.TOK_PER_WARP + self.ctx.rgroup * MFMA_MNK + self.ctx.lane16 - ) * self.traits.f32 + ) + in_context = tile_tok < tile_valid + k_scale_scalar = in_context.select(k_scale_scalar, fx.Float32(0.0)) + v_scale_scalar = in_context.select(v_scale_scalar, fx.Float32(0.0)) + slot = tile_tok * self.traits.f32 self.lds.store( self.traits.sKScale_off + buf_off + slot, fx.Float32, @@ -681,6 +701,73 @@ def load_v(self, phys_row, vh): fx.rocdl.sched_vmem(len(ops) // 2) return ops # NVOPS i64, the 64-token contiguous run for this head + def v_tail_masks(self, tile_valid): + """Byte masks keeping only this lane's V tokens below the context length. + + Operand i holds tokens rgroup * TOK_PER_WARP + i * (8 // element bytes) + onward, in load_v's order. + """ + remaining = tile_valid - self.ctx.rgroup * self.traits.TOK_PER_WARP + bits_per_tok = 8 * self.traits.QP_ELEM_BYTES + if const_expr(self.traits.is_bf16_kv and self.traits.block_size >= 64): + # The per-word form below is fewer VALU ops, but these kernels + # already sit at the register limit and it spills inside the tile + # loop there. + tok_per_op = 8 // self.traits.QP_ELEM_BYTES + masks = [] + for i in range_constexpr(self.traits.NVOPS): + n = remaining - i * tok_per_op + n = (n > 0).select(n, 0) + full = n >= tok_per_op + shift = fx.Int64((n * bits_per_tok) & 63) + masks.append(full.select(fx.Int64(-1), (fx.Int64(1) << shift) - 1)) + return masks + valid_bits = remaining * bits_per_tok + words = [] + for d in range_constexpr(2 * self.traits.NVOPS): + # A 64-bit shift of the low word by 0..32 bits covers every 32-bit + # mask from all ones to zero. + shift = 32 * (d + 1) - valid_bits + shift = (shift > 0).select(shift, 0) + shift = (shift < 32).select(shift, 32) + words.append(fx.Int64(0xFFFFFFFF) >> fx.Int64(shift)) + return [ + words[2 * i] | (words[2 * i + 1] << 32) + for i in range_constexpr(self.traits.NVOPS) + ] + + @flyc.jit + def tail_masks(self, tile_valid): + # Only a sequence's last tile holds tokens at or past its context. + masks = fx.full(self.traits.NVOPS, -1, fx.Int64) + if tile_valid < self.traits.TILE_TOK: + masks = fx.Vector.from_elements( + self.v_tail_masks(tile_valid), dtype=fx.Int64 + ) + return masks + + @flyc.jit + def _mask_v_chunk(self, v, masks, tile_valid): + if tile_valid < self.traits.TILE_TOK: + v = v & masks + return v + + def mask_v_chunk(self, v_ops, masks, tile_valid): + """Zero one head chunk's V tokens at or past the context length before PV. + + Their probabilities are exactly zero, but those rows may be pinned + block-0 pages or unwritten tails holding NaN, and 0 * NaN is NaN in PV. + """ + v = self._mask_v_chunk( + fx.Vector.from_elements( + [v_ops[i] for i in range_constexpr(self.traits.NVOPS)], + dtype=fx.Int64, + ), + masks, + tile_valid, + ) + return [v[i] for i in range_constexpr(self.traits.NVOPS)] + @flyc.jit def stage_page_ids(self): """Stage this task's page ids in LDS so the tile loop never reads the table.""" @@ -713,7 +800,9 @@ def prefetch_first(self): # writes does not also drain the first K prefetch. phys_vec0 = self.fetch_v_pages(self.ctx.part_start) self.stage_scales( - phys_vec0, self.scale_buffer_offset(fx.Int32(self.ctx.part_start)) + phys_vec0, + self.scale_buffer_offset(fx.Int32(self.ctx.part_start)), + self.ctx.part_start, ) k_pf0 = self.load_k_from_pages(phys_vec0) else: @@ -727,6 +816,7 @@ def prefetch_first(self): self.stage_scales( phys_vec0, self.scale_buffer_offset(fx.Int32(self.ctx.part_start)), + self.ctx.part_start, ) elif self.ctx.lane == 0: self.lds.store( diff --git a/aiter/ops/flydsl/kernels/pa_decode/pipeline.py b/aiter/ops/flydsl/kernels/pa_decode/pipeline.py index 250fa974df9..9a4cca12f9d 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/pipeline.py +++ b/aiter/ops/flydsl/kernels/pa_decode/pipeline.py @@ -248,12 +248,14 @@ def multi_query_tile( if const_expr(self.traits.MTP4_FUSED): # Defer K until P packing releases saved score registers. phys_vec1 = self.kv.fetch_v_pages(tt1) - self.kv.stage_scales(phys_vec1, self.kv.scale_buffer_offset(tt1)) + self.kv.stage_scales(phys_vec1, self.kv.scale_buffer_offset(tt1), tt1) else: k_next, phys_vec1 = self.kv.load_k(tt1) self.kv.fetch_v_pages(tt1) if const_expr(self.traits.per_token_kv): - self.kv.stage_scales(phys_vec1, self.kv.scale_buffer_offset(tt1)) + self.kv.stage_scales( + phys_vec1, self.kv.scale_buffer_offset(tt1), tt1 + ) next_state[self.traits.K_SLOT] = k_next gpu.barrier() @@ -298,6 +300,7 @@ def multi_query_tile( gpu.barrier() + v_masks = self.kv.tail_masks(tile_valid) v_next_chunks = [] for m in range_constexpr(self.traits.M_TILES): p_base = self.traits.sP_off + m * MFMA_MNK * self.traits.SP_ROW_BYTES @@ -335,6 +338,10 @@ def multi_query_tile( ) corr_b = fx.full(self.traits.OP_ELEMS, fx.Float32(corr_reg), fx.Float32) for vh in range_constexpr(self.traits.VHE_CHUNKS): + if const_expr(m == 0): + v_vh_shared[vh] = self.kv.mask_v_chunk( + v_vh_shared[vh], v_masks, tile_valid + ) v_vh = v_vh_shared[vh] op = self.gemm.pv(v_vh, p_ops) * fx.full( self.traits.OP_ELEMS, v_max_scaled, fx.Float32 @@ -435,7 +442,13 @@ def multi_query_tile( ) corr_b = fx.full(self.traits.OP_ELEMS, fx.Float32(corr_reg), fx.Float32) + if const_expr(m == 0): + v_masks = self.kv.tail_masks(tile_valid) for vh in range_constexpr(self.traits.VHE_CHUNKS): + if const_expr(m == 0): + v_vh_shared[vh] = self.kv.mask_v_chunk( + v_vh_shared[vh], v_masks, tile_valid + ) v_vh = v_vh_shared[vh] op = self.gemm.pv(v_vh, p_ops) if const_expr(self.traits.per_token_kv): @@ -506,14 +519,18 @@ def single_query_tile( # page16 defers K until the P barrier. phys_vec1 = self.kv.fetch_v_pages(tt1) if const_expr(self.traits.per_token_kv): - self.kv.stage_scales(phys_vec1, self.kv.scale_buffer_offset(tt1)) + self.kv.stage_scales( + phys_vec1, self.kv.scale_buffer_offset(tt1), tt1 + ) if const_expr(not self.traits.PAGE16_VPIPE): k_next = self.kv.load_k_from_pages(phys_vec1) else: k_next, phys_vec1 = self.kv.load_k(tt1) self.kv.fetch_v_pages(tt1) if const_expr(self.traits.per_token_kv): - self.kv.stage_scales(phys_vec1, self.kv.scale_buffer_offset(tt1)) + self.kv.stage_scales( + phys_vec1, self.kv.scale_buffer_offset(tt1), tt1 + ) if const_expr(self.traits.SCALES_BEFORE_CURRENT_V): # Issue V after scale staging so scale waits do not drain V. v_vh_shared = [ @@ -597,10 +614,11 @@ def single_query_tile( self.kv.load_v(v_page_cur, vh) for vh in range_constexpr(self.traits.VHE_CHUNKS) ] + v_masks = self.kv.tail_masks(tile_valid) # vh outer, m inner: each M-tile is a separate accumulator chain on the # same V operand. for vh in range_constexpr(self.traits.VHE_CHUNKS): - v_vh = v_vh_batch[vh] + v_vh = self.kv.mask_v_chunk(v_vh_batch[vh], v_masks, tile_valid) for m in range_constexpr(self.traits.M_TILES): op = self.gemm.pv(v_vh, p_ops[m]) if const_expr(self.traits.per_token_kv): diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 74f174fe892..1e008ea0a5f 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -173,9 +173,12 @@ def compile_pa_decode_tile( are absent and query/output strides match. Other tasks write packed partials for the shared reducer, which applies sinks once. - Masked V bytes must remain finite because ``0 * NaN == NaN`` in PV MFMA. - Pages past the sequence are pinned to block 0; callers must leave the - unwritten tail of the last owned page finite. ``wide_kv_addressing`` protects + Pages past the sequence are pinned to block 0 and the unwritten tail of the + last owned page is read as is. On gfx942/gfx950 a sequence's last tile + zeroes V bytes and K/V scales at or past its context before PV, so those + rows may hold anything, including NaN (``0 * NaN == NaN`` in PV MFMA). + The gfx1250 pipelines still require those V bytes to remain finite. + ``wide_kv_addressing`` protects offsets at 2 GiB; ``kv_buffer_u32`` proves both FP8 caches are below 4 GiB and permits unsigned buffer offsets instead of i64. """ diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 7fae13c2dcd..7fc0e13ad57 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -746,3 +746,128 @@ def check(): output.fill_(float("nan")) graph.replay() check() + + +def _poison_unowned(cache, scale, owned, fill): + """Overwrite K/V rows no sequence owns (and their per-token scales).""" + rows = ~owned[:, None, :].expand(-1, cache.shape[1], -1) + if cache.dtype == torch.bfloat16: + bits = {"zero": 0, "nan": 0x7FC0, "inf": 0x7F80}[fill] + cache.view(torch.int16)[rows] = bits + else: + # E4M3 has no infinity; its largest finite magnitude stands in. + nan = 0x80 if cache.dtype == torch.float8_e4m3fnuz else 0x7F + bits = {"zero": 0, "nan": nan, "inf": 0x7E}[fill] + cache.view(torch.uint8)[rows] = bits + if scale is not None and scale.numel() > 1: + scale.view(scale.shape[:3])[rows] = {"zero": 0.0, "nan": float("nan")}.get( + fill, float("inf") + ) + + +@pytest.mark.parametrize("fill", ["nan", "inf"]) +@pytest.mark.parametrize( + "query_length,kv_heads,group_size,head_dim", + [(1, 1, 8, 128), (4, 2, 8, 256), (4, 1, 16, 128)], +) +@pytest.mark.parametrize("block_size", [16, 64, 128]) +@pytest.mark.parametrize( + "kv", ["fp8-per-tensor", "fp8-per-tensor-plain-v", "fp8-per-token", "bf16"] +) +def test_pa_decode_unowned_kv_rows( + device, kv, block_size, query_length, kv_heads, group_size, head_dim, fill +): + """Rows no sequence owns contribute exactly zero, even when NaN or Inf. + + Padded block-table lanes read block 0 and the last page's tail past the + context is read as is. Both may hold NaN (a recycled page, a checkpoint + image); P is zero there but 0 * NaN is NaN in PV. + """ + if get_gfx_runtime() == "gfx1250": + pytest.skip("gfx1250 pipelines still require finite unowned V rows") + pa = importlib.import_module("aiter.ops.flydsl.pa_decode") + torch.manual_seed(0) + lengths = tuple(max(n, query_length) for n in (1, 17, 255, 257, 1000)) + query, key, value, table, context = _make_inputs( + query_length, + kv_heads, + group_size, + head_dim, + block_size, + torch.bfloat16, + "random", + device, + lengths, + ) + # Page 0 is owned by no sequence; padded block-table lanes point at it. + spare = torch.zeros_like(key[:1]) + key, value = torch.cat([spare, key]), torch.cat([spare, value]) + table = torch.nn.functional.pad(table + 1, (0, 2)) + owned = torch.zeros(key.shape[0], block_size, dtype=torch.bool, device=device) + for seq, length in enumerate(lengths): + tokens = torch.arange(length, device=device) + owned[table[seq, tokens // block_size].long(), tokens % block_size] = True + + key_scale = value_scale = None + vector = 8 + if kv.startswith("fp8"): + compute_type = ( + torch.float8_e4m3fn + if get_gfx_runtime() == "gfx950" + else torch.float8_e4m3fnuz + ) + vector = 16 + quantize = pertoken_quant if kv == "fp8-per-token" else per_tensor_quant + key, key_scale = quantize(key, quant_dtype=compute_type) + value, value_scale = quantize(value, quant_dtype=compute_type) + logical_key = key.float() * (1.0 if key_scale is None else key_scale) + logical_value = value.float() * (1.0 if value_scale is None else value_scale) + reference = _reference( + query, logical_key, logical_value, table, context, query_length, 0, None + ) + plan = pa.plan_pa_decode(context, kv_heads, query_length=query_length) + + def run(fill): + k, v = key.clone(), value.clone() + ks = None if key_scale is None else key_scale.clone() + vs = None if value_scale is None else value_scale.clone() + _poison_unowned(k, ks, owned, fill) + _poison_unowned(v, vs, owned, fill) + pages = k.shape[0] + key_cache = ( + k.reshape(pages, kv_heads, block_size, head_dim // vector, vector) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + if kv == "fp8-per-tensor-plain-v": + value_cache = v.permute(0, 1, 3, 2).contiguous() + else: + value_cache = ( + v.reshape(pages, kv_heads, block_size // vector, vector, head_dim) + .permute(0, 1, 2, 4, 3) + .contiguous() + ) + output = torch.full_like(query, float("nan")) + pa.pa_decode( + output, + query, + key_cache, + value_cache, + context, + table, + head_dim**-0.5, + query_length, + compute_type=key_cache.dtype, + key_scale=ks, + value_scale=vs, + work_plan=plan, + max_context_length=max(lengths), + ) + return output + + clean, poisoned = run("zero"), run(fill) + assert torch.isfinite(poisoned).all() + # Short FP8 contexts round P coarsely; exactness is the bitwise check. + tol = 5e-3 if kv == "bf16" else 2.5e-2 + torch.testing.assert_close(poisoned.float(), reference, atol=tol, rtol=tol) + assert torch.equal(poisoned, clean) From 795de30246865cbe301da3e772b7859f0873bec3 Mon Sep 17 00:00:00 2001 From: fsx950223 Date: Thu, 29 Oct 2026 07:36:27 +0000 Subject: [PATCH 2/2] [FlyDSL] pa_decode: add mask_unowned_kv option to toggle V masking mask_unowned_kv (default True) keeps the boundary-tile V/scale masking introduced in this PR. False compiles it out, restoring the previous kernel and the caller contract that unowned V bytes and scales are finite. The flag is part of the schedule cache key; gfx1250 ignores it. Co-Authored-By: Claude Opus 5.5 --- aiter/ops/flydsl/kernels/pa_decode/op_lds.py | 20 ++++++++++++++++---- aiter/ops/flydsl/kernels/pa_decode/traits.py | 5 +++++ aiter/ops/flydsl/kernels/pa_decode_kernel.py | 12 ++++++++---- aiter/ops/flydsl/pa_decode.py | 9 +++++++++ op_tests/test_flydsl_pa_decode.py | 6 +++++- 5 files changed, 43 insertions(+), 9 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py b/aiter/ops/flydsl/kernels/pa_decode/op_lds.py index 6b326acf45a..463925e5d5a 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py +++ b/aiter/ops/flydsl/kernels/pa_decode/op_lds.py @@ -504,6 +504,8 @@ def _context_scales(self, scales, remaining): P is zero there, but a NaN or infinite V scale still makes P * scale NaN. ``remaining`` counts valid tokens from the vector's first element. """ + if const_expr(not self.traits.mask_unowned_kv): + return scales zero = fx.Float32(0.0) return fx.Vector.from_elements( [ @@ -557,9 +559,10 @@ def stage_scales(self, phys_vec, buf_off, tt): + self.ctx.rgroup * MFMA_MNK + self.ctx.lane16 ) - in_context = tile_tok < tile_valid - k_scale_scalar = in_context.select(k_scale_scalar, fx.Float32(0.0)) - v_scale_scalar = in_context.select(v_scale_scalar, fx.Float32(0.0)) + if const_expr(self.traits.mask_unowned_kv): + in_context = tile_tok < tile_valid + k_scale_scalar = in_context.select(k_scale_scalar, fx.Float32(0.0)) + v_scale_scalar = in_context.select(v_scale_scalar, fx.Float32(0.0)) slot = tile_tok * self.traits.f32 self.lds.store( self.traits.sKScale_off + buf_off + slot, @@ -736,8 +739,14 @@ def v_tail_masks(self, tile_valid): for i in range_constexpr(self.traits.NVOPS) ] - @flyc.jit def tail_masks(self, tile_valid): + """V byte masks for mask_v_chunk; None when masking is disabled.""" + if const_expr(not self.traits.mask_unowned_kv): + return None + return self._tail_masks(tile_valid) + + @flyc.jit + def _tail_masks(self, tile_valid): # Only a sequence's last tile holds tokens at or past its context. masks = fx.full(self.traits.NVOPS, -1, fx.Int64) if tile_valid < self.traits.TILE_TOK: @@ -757,7 +766,10 @@ def mask_v_chunk(self, v_ops, masks, tile_valid): Their probabilities are exactly zero, but those rows may be pinned block-0 pages or unwritten tails holding NaN, and 0 * NaN is NaN in PV. + Disabled by ``mask_unowned_kv=False``, which returns ``v_ops`` as is. """ + if const_expr(not self.traits.mask_unowned_kv): + return v_ops v = self._mask_v_chunk( fx.Vector.from_elements( [v_ops[i] for i in range_constexpr(self.traits.NVOPS)], diff --git a/aiter/ops/flydsl/kernels/pa_decode/traits.py b/aiter/ops/flydsl/kernels/pa_decode/traits.py index de94a95d83f..ce14f320779 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/traits.py +++ b/aiter/ops/flydsl/kernels/pa_decode/traits.py @@ -127,6 +127,7 @@ def select( max_context_length: int | None = None, sliding_window: int = 0, trans_k: bool = True, + mask_unowned_kv: bool = True, ) -> PaDecodeSchedule: """Choose the existing policy from host-known grid and query bounds.""" if isinstance(work_capacity, bool) or not isinstance(work_capacity, int): @@ -345,6 +346,7 @@ def select( batch_first_plan_grid, sliding_window, kv_dtype, + bool(mask_unowned_kv), D256_PIPE, stage_pages, ) @@ -374,6 +376,7 @@ class PaDecodeTraits: query_length: int trans_v: bool wide_kv_addressing: bool + mask_unowned_kv: bool query_splits: int sliding_window: int single_tile_plan: bool @@ -482,6 +485,7 @@ def create(cls, schedule: PaDecodeSchedule) -> PaDecodeTraits: batch_first_plan_grid, sliding_window, kv_dtype, + mask_unowned_kv, D256_PIPE, stage_page_capacity, ) = cache_key @@ -681,6 +685,7 @@ def create(cls, schedule: PaDecodeSchedule) -> PaDecodeTraits: query_length=query_length, trans_v=trans_v, wide_kv_addressing=wide_kv_addressing, + mask_unowned_kv=mask_unowned_kv, query_splits=query_splits, sliding_window=sliding_window, single_tile_plan=single_tile_plan, diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 1e008ea0a5f..986f37ae00c 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/aiter/ops/flydsl/kernels/pa_decode_kernel.py @@ -140,6 +140,7 @@ def compile_pa_decode_tile( sliding_window: int = 0, direct_output: bool = False, trans_k: bool = True, + mask_unowned_kv: bool = True, ): """Select and cache a PA-decode kernel and launch wrapper. @@ -174,10 +175,12 @@ def compile_pa_decode_tile( for the shared reducer, which applies sinks once. Pages past the sequence are pinned to block 0 and the unwritten tail of the - last owned page is read as is. On gfx942/gfx950 a sequence's last tile - zeroes V bytes and K/V scales at or past its context before PV, so those - rows may hold anything, including NaN (``0 * NaN == NaN`` in PV MFMA). - The gfx1250 pipelines still require those V bytes to remain finite. + last owned page is read as is. On gfx942/gfx950 with ``mask_unowned_kv`` + (the default), a sequence's last tile zeroes V bytes and K/V scales at or + past its context before PV, so those rows may hold anything, including NaN + (``0 * NaN == NaN`` in PV MFMA). ``mask_unowned_kv=False`` skips that + masking and restores the caller contract that those V bytes and scales + remain finite. The gfx1250 pipelines ignore the flag and always require it. ``wide_kv_addressing`` protects offsets at 2 GiB; ``kv_buffer_u32`` proves both FP8 caches are below 4 GiB and permits unsigned buffer offsets instead of i64. @@ -241,6 +244,7 @@ def compile_pa_decode_tile( max_context_length=max_context_length, sliding_window=sliding_window, trans_k=trans_k, + mask_unowned_kv=mask_unowned_kv, ) cached = _PA_DECODE_TILE_CACHE.get(schedule.cache_key) if cached is not None: diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 4e33a3e3190..ac7de520e5f 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/aiter/ops/flydsl/pa_decode.py @@ -193,6 +193,7 @@ def pa_decode( sliding_window: int = 0, work_plan: PADecodePlan, max_context_length: int | None = None, + mask_unowned_kv: bool = True, ) -> None: """Decode FP8 or BF16 K/V using 256-token compute tiles. @@ -250,6 +251,13 @@ def pa_decode( ``max_context_length``. A smaller value still covers the real context: page ids outside the staged window are read from the block table. + ``mask_unowned_kv`` (default True) zeroes, on gfx942/gfx950, V bytes and + K/V scales of tokens at or past each sequence's context length before PV, + so pinned block-0 pages and the unwritten tail of the last page may hold + any bytes, including NaN. False skips that masking (slightly faster on + tiny boundary-tile-only calls); the caller must then keep those bytes and + scales finite. gfx1250 ignores it and always requires finite bytes. + GPU lengths and table entries are not checked. Callers must ensure 0 <= context_lengths[i] <= block_tables.shape[1] * block_size and used block indices in [0, min(key_cache.shape[0], value_cache.shape[0])). @@ -644,6 +652,7 @@ def normalize_scale(scale, name): sliding_window=sliding_window, direct_output=direct_output, trans_k=trans_k, + mask_unowned_kv=mask_unowned_kv, ) total_rows = query_length * query_group_size diff --git a/op_tests/test_flydsl_pa_decode.py b/op_tests/test_flydsl_pa_decode.py index 7fc0e13ad57..3a6a9f65166 100644 --- a/op_tests/test_flydsl_pa_decode.py +++ b/op_tests/test_flydsl_pa_decode.py @@ -782,6 +782,7 @@ def test_pa_decode_unowned_kv_rows( Padded block-table lanes read block 0 and the last page's tail past the context is read as is. Both may hold NaN (a recycled page, a checkpoint image); P is zero there but 0 * NaN is NaN in PV. + ``mask_unowned_kv=False`` skips the masking and matches it on clean input. """ if get_gfx_runtime() == "gfx1250": pytest.skip("gfx1250 pipelines still require finite unowned V rows") @@ -827,7 +828,7 @@ def test_pa_decode_unowned_kv_rows( ) plan = pa.plan_pa_decode(context, kv_heads, query_length=query_length) - def run(fill): + def run(fill, mask_unowned_kv=True): k, v = key.clone(), value.clone() ks = None if key_scale is None else key_scale.clone() vs = None if value_scale is None else value_scale.clone() @@ -862,6 +863,7 @@ def run(fill): value_scale=vs, work_plan=plan, max_context_length=max(lengths), + mask_unowned_kv=mask_unowned_kv, ) return output @@ -871,3 +873,5 @@ def run(fill): tol = 5e-3 if kv == "bf16" else 2.5e-2 torch.testing.assert_close(poisoned.float(), reference, atol=tol, rtol=tol) assert torch.equal(poisoned, clean) + if fill == "nan": + assert torch.equal(run("zero", mask_unowned_kv=False), clean)