[FlyDSL] pa_decode: zero V rows past the context before PV

aiter/ops/flydsl/kernels/pa_decode/op_lds.py

@@ -498,7 +498,25 @@ 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.
+        """
+        if const_expr(not self.traits.mask_unowned_kv):
+            return scales
+        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 +532,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 +554,16 @@ 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
+            )
+            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,
                 fx.Float32,
@@ -681,6 +704,82 @@ 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)
+        ]
+
+    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:
+            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.
+        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)],
+                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 +812,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 +828,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(

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

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,

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.
 
@@ -173,9 +174,14 @@ 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 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.
     """
@@ -238,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:

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