[FlyDSL] fp8_mqa_logits: drop the host row-pad copies

aiter/ops/flydsl/fp8_mqa_logits_kernels.py

@@ -9,8 +9,8 @@
     (``_auto_variant`` / ``KERNEL_VARIANTS`` / ``DEFAULT_VARIANT``).
   - A build cache keyed by shape/variant/dtype-conversion flags
     (``compile_fp8_mqa_logits``).
-  - Host-side seq_len padding, output-column alignment, and the KV-column
-    split (``grid.y``) heuristic that fills the device for small-M shapes.
+  - Output-column alignment, and the KV-column split (``grid.y``) heuristic
+    that fills the device for small-M shapes. Query rows are not padded.
 """
 
 from __future__ import annotations
@@ -160,7 +160,7 @@ def _max_blocks_per_cu(launcher, args):
 
 
 def _auto_num_splits(
-    seq_len_padded: int,
+    seq_len: int,
     seq_len_kv: int,
     rows_per_block: int,
     block_kv: int,
@@ -180,7 +180,7 @@ def _auto_num_splits(
     the compiled kernel (its register and LDS use), not only on the device.
     """
     pol = _split_policy()
-    grid_x = seq_len_padded // rows_per_block
+    grid_x = (seq_len + rows_per_block - 1) // rows_per_block
     if grid_x == 0 or seq_len_kv < pol.min_seq_len_kv:
         return 1
     max_splits = max(1, (seq_len_kv // block_kv) // pol.min_tiles_per_split)
@@ -382,10 +382,10 @@ def _register_variants(entries):
 
 
 def _parse_variant(tag):
-    """(block_kv, rows_per_block_effective) for host-side padding and splitting.
+    """(block_kv, rows_per_block_effective) for the row grid and KV splits.
 
     For ``_lds`` variants the WPB waves partition rows within one shared KV
-    tile, so a block owns RPB*WPB rows and seq_len must be padded to that.
+    tile, so a block owns RPB*WPB rows. A short tail is masked in the kernel.
     """
     m = _TAG_RE.match(tag)
     if m is None:
@@ -399,8 +399,7 @@ def _auto_variant(seq_len, seq_len_kv, num_heads):
     """Pick a variant from the problem shape.
 
     gfx942: RPB from ``seq_len * seq_len_kv`` thresholds; in the middle band an
-        odd ``seq_len`` drops RPB to 1, since padding costs a fixed host-side
-        overhead that only the largest shapes outgrow. WPB=2 packs more column
+        odd ``seq_len`` stays at RPB 1. WPB=2 packs more column
         tiles per wave when M and N are both large, else WPB=4 for more
         wavefronts on small-M / short-window shapes.
 
@@ -614,26 +613,11 @@ def flydsl_fp8_mqa_logits(
         clean_logits=bool(clean_logits),
     )
 
-    # The kernels require seq_len padded to a multiple of the rows a block owns,
-    # so every block owns exactly that many. Padded rows get empty windows
-    # (start == end == 0) so the kernel writes nothing for them; the output is
-    # sliced back to the original seq_len after the launch.
-    seq_len_padded = (
-        (seq_len + _ROWS_PER_BLOCK - 1) // _ROWS_PER_BLOCK
-    ) * _ROWS_PER_BLOCK
-    if seq_len_padded != seq_len:
-        pad = seq_len_padded - seq_len
-        Q = torch.cat([Q, Q.new_zeros((pad, num_heads, head_size))], dim=0)
-        weights = torch.cat([weights, weights.new_zeros((pad, num_heads))], dim=0)
-        cu_starts = torch.cat([cu_starts, cu_starts.new_zeros(pad)], dim=0)
-        cu_ends = torch.cat([cu_ends, cu_ends.new_zeros(pad)], dim=0)
-
-    # No torch.full even when clean_logits: the kernel now writes -inf itself,
-    # at exactly the out-of-window positions it would otherwise skip.
+    # No torch.full even when clean_logits: the kernel writes -inf itself.
     aligned_size = 256
     seq_len_kv_aligned = (seq_len_kv + aligned_size - 1) // aligned_size * aligned_size
     logits = torch.empty(
-        (seq_len_padded, seq_len_kv_aligned),
+        (seq_len, seq_len_kv_aligned),
         dtype=torch.float32,
         device=Q.device,
     )[:, :seq_len_kv]
@@ -642,7 +626,7 @@ def flydsl_fp8_mqa_logits(
         stream = torch.cuda.current_stream()
 
     tensors = (Q, KV, kv_scales, weights, cu_starts, cu_ends, logits)
-    scalars = (int(seq_len_padded), int(seq_len_kv), int(logits.stride(0)))
+    scalars = (int(seq_len), int(seq_len_kv), int(logits.stride(0)))
 
     with torch.cuda.device(Q.device.index):
         blocks_per_cu = None
@@ -651,7 +635,7 @@ def flydsl_fp8_mqa_logits(
                 launcher, (*tensors, *scalars, 1, stream)
             )
         num_splits = _auto_num_splits(
-            seq_len_padded,
+            seq_len,
             seq_len_kv,
             _ROWS_PER_BLOCK,
             _BKV,

aiter/ops/flydsl/kernels/mqa_logits/fp8_mqa_logits.py

@@ -124,6 +124,7 @@ def _emit_row_neg_inf_fill(
     starts,  # list[fx.Int32]: max(cu_starts, 0),        parallel to rows
     ends,  # list[fx.Int32]: min(cu_ends, seq_len_kv), parallel to rows
     seq_len_kv,  # fx.Int32
+    seq_len,  # fx.Int32
     by_i32,  # fx.Int32: block_idx.y
     num_splits,  # fx.Int32: grid.y (>= 1)
     fill_range,  # (out_row_t, lo, hi) -> None: thread-strided -inf fill
@@ -178,10 +179,16 @@ def _split(lo, hi):
         b_lo = fx.min(lo + by_i32 * chunk, hi)
         return b_lo, fx.min(b_lo + chunk, hi)
 
+    seq_len_m_1 = seq_len - fx.Int32(1)
     for j in range_constexpr(len(rows)):
-        out_row_t = _make_out_row_t(logits, stride_i64, rows[j])
+        # An empty window fills the whole row, so a past-the-end slot must
+        # collapse both ranges instead. The descriptor stays on a real row.
+        in_rows = rows[j] < seq_len
+        out_row_t = _make_out_row_t(logits, stride_i64, fx.min(rows[j], seq_len_m_1))
         s = fx.min(starts[j], slk)
         e = fx.max(ends[j], s)
+        s = in_rows.select(s, fx.Int32(0))
+        e = in_rows.select(e, slk)
         for lo_i32, hi_i32 in ((fx.Int32(0), s), (e, slk)):
             b_lo, b_hi = _split(lo_i32, hi_i32)
             fill_range(out_row_t, b_lo, b_hi)
@@ -362,10 +369,8 @@ def _build_kernel_mfma_r_w(
         within each BKV tile.
       * A-operand (Q) layout and head-reduce are per-lane within the wave (width 64).
 
-    Grid: ``(ceil(seq_len / RPB), num_splits, 1)``.  The host pads ``seq_len`` to
-    a multiple of ``RPB`` (every block owns exactly ``RPB`` rows) and may split
-    each row's KV window across ``grid.y`` blocks when the row grid alone is too
-    small to fill the device (see ``flydsl_fp8_mqa_logits``).
+    Grid: ``(ceil(seq_len / RPB), num_splits, 1)``. The last block may be short.
+    The host may split each row's KV window across ``grid.y``.
     """
     H = num_heads
     D = head_size
@@ -425,7 +430,7 @@ def kernel(
         cu_starts: fx.Tensor,  # [seq_len]             i32
         cu_ends: fx.Tensor,  # [seq_len]             i32
         logits: fx.Tensor,  # [seq_len, seq_len_kv] f32
-        seq_len: fx.Int32,  # padded to a multiple of RPB
+        seq_len: fx.Int32,
         seq_len_kv: fx.Int32,
         stride_logits_s: fx.Int32,
         num_splits: fx.Int32,  # grid.y KV-column splits (1 == no split)
@@ -512,22 +517,29 @@ def _fix_i32(src):
 
         # ---- Preload window bounds, Q frags, and weights for all RPB rows ----
         # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q.
+        # Past-the-end slots reuse the last real row's addresses and take an
+        # empty window, so they are not stored.
         starts = [None] * RPB
         ends = [None] * RPB
         a_packs = [None] * RPB
         w_frag = [None] * RPB
+        seq_len_m_1 = seq_len - fx.Int32(1)
 
         for j in range_constexpr(RPB):
             row = r0 + fx.Int32(j)
-            starts[j] = fx.max(fx.Int32(cs_t[row]), fx.Int32(0))
-            ends[j] = fx.min(fx.Int32(ce_t[row]), seq_len_kv)
+            row_ld = fx.min(row, seq_len_m_1)
+            in_rows = row < seq_len
+            ss = fx.max(fx.Int32(cs_t[row_ld]), fx.Int32(0))
+            ee = fx.min(fx.Int32(ce_t[row_ld]), seq_len_kv)
+            starts[j] = in_rows.select(ss, fx.Int32(0))
+            ends[j] = in_rows.select(ee, fx.Int32(0))
 
             # lane -> Q[row, h = mi*MFMA_M + lane%MFMA_N,
             #            d = kk*MFMA_K + (lane//MFMA_N)*8 + 0..7]
             row_a = [[None] * K_STEPS for _ in range_constexpr(M_TILES)]
             for mi in range_constexpr(M_TILES):
                 h_a = fx.Int32(mi * MFMA_M) + lane_mod_N
-                row_h = row * fx.Int32(H) + h_a
+                row_h = row_ld * fx.Int32(H) + h_a
                 base_a = row_h * fx.Int32(D)
                 for kk in range_constexpr(K_STEPS):
                     row_a[mi][kk] = _load_frag(
@@ -536,7 +548,7 @@ def _fix_i32(src):
             a_packs[j] = row_a
 
             w_frag[j] = _load_row_weights(
-                weights, H, cp_4xfp32, tc_c_w, M_TILES, MFMA_N, row
+                weights, H, cp_4xfp32, tc_c_w, M_TILES, MFMA_N, row_ld
             )
 
         # ---- Union window across all RPB rows ----
@@ -649,6 +661,7 @@ def _fill_range(out_row_t, lo_i32, hi_i32):
                 starts=starts,
                 ends=ends,
                 seq_len_kv=seq_len_kv,
+                seq_len=seq_len,
                 by_i32=by,
                 num_splits=num_splits,
                 fill_range=_fill_range,
@@ -825,7 +838,7 @@ def kernel(
         cu_starts: fx.Tensor,
         cu_ends: fx.Tensor,
         logits: fx.Tensor,
-        seq_len: fx.Int32,  # padded to a multiple of ROWS_PER_BLOCK
+        seq_len: fx.Int32,
         seq_len_kv: fx.Int32,
         stride_logits_s: fx.Int32,
         num_splits: fx.Int32,
@@ -929,12 +942,17 @@ def _dma_kv_tile_to_lds(slot_byte_i32, col0_i32):
         ends = [None] * RPW
         a_packs = [None] * RPW
         w_frag = [None] * RPW
+        seq_len_m_1 = seq_len - fx.Int32(1)
 
         # Loop over the rows owned by this wave.
         for j in range_constexpr(RPW):
             row = wave_row0 + fx.Int32(j)
-            starts[j] = fx.max(fx.Int32(cs_t[row]), fx.Int32(0))
-            ends[j] = fx.min(fx.Int32(ce_t[row]), seq_len_kv)
+            row_ld = fx.min(row, seq_len_m_1)
+            in_rows = row < seq_len
+            ss = fx.max(fx.Int32(cs_t[row_ld]), fx.Int32(0))
+            ee = fx.min(fx.Int32(ce_t[row_ld]), seq_len_kv)
+            starts[j] = in_rows.select(ss, fx.Int32(0))
+            ends[j] = in_rows.select(ee, fx.Int32(0))
 
             # Load A-frags:
             # Q[
@@ -945,7 +963,7 @@ def _dma_kv_tile_to_lds(slot_byte_i32, col0_i32):
             row_a_frag = [[None] * K_STEPS for _ in range_constexpr(M_TILES)]
             for mi in range_constexpr(M_TILES):
                 h_a = fx.Int32(mi * mfma.MFMA_M) + lane_mod_N
-                row_h = h_a + row * fx.Int32(H)
+                row_h = h_a + row_ld * fx.Int32(H)
                 base_a = row_h * fx.Int32(D)
                 for kk in range_constexpr(K_STEPS):
                     row_a_frag[mi][kk] = mfma.make_frag(
@@ -957,7 +975,7 @@ def _dma_kv_tile_to_lds(slot_byte_i32, col0_i32):
             a_packs[j] = row_a_frag
 
             w_frag[j] = _load_row_weights(
-                weights, H, cp_4xfp32, tc_c_w, M_TILES, mfma.MFMA_N, row
+                weights, H, cp_4xfp32, tc_c_w, M_TILES, mfma.MFMA_N, row_ld
             )
 
         # ---- Union KV window across all block rows (all waves cooperate) ----
@@ -970,8 +988,12 @@ def _dma_kv_tile_to_lds(slot_byte_i32, col0_i32):
         #   u_end = max(cu_ends[rows]).
         for jj in range_constexpr(ROWS_PER_BLOCK):
             rr = block_row0 + fx.Int32(jj)
-            ss = fx.max(fx.Int32(cs_t[rr]), fx.Int32(0))
-            ee = fx.min(fx.Int32(ce_t[rr]), seq_len_kv)
+            rr_ld = fx.min(rr, seq_len_m_1)
+            in_rows = rr < seq_len
+            ss = fx.max(fx.Int32(cs_t[rr_ld]), fx.Int32(0))
+            ee = fx.min(fx.Int32(ce_t[rr_ld]), seq_len_kv)
+            ss = in_rows.select(ss, fx.Int32(0))
+            ee = in_rows.select(ee, fx.Int32(0))
             if jj == 0:
                 u_start = ss
                 u_end = ee
@@ -1135,6 +1157,7 @@ def _fill_range(out_row_t, lo_i32, hi_i32):
                 starts=starts,
                 ends=ends,
                 seq_len_kv=seq_len_kv,
+                seq_len=seq_len,
                 by_i32=block_y,
                 num_splits=num_splits,
                 fill_range=_fill_range,

op_tests/test_flydsl_fp8_mqa_logits.py

@@ -416,7 +416,7 @@ def verify_fp8_mqa_logits(
     (65, 8192): "mfma_r1_w4",  # odd, below 2**21: steps down
     (254, 8192): "mfma_r2_w4",
     (256, 8192): "mfma_r4_w4",  # 2**21
-    (257, 8192): "mfma_r4_w4",  # odd, from 2**21 up: pads
+    (257, 8192): "mfma_r4_w4",  # odd, from 2**21 up: short tail tile
     (1024, 131072): "mfma_r4_w4",  # long-context indexer prefill
     (2048, 8192): "mfma_r4_w2",
 }
@@ -780,8 +780,8 @@ def main():
             # enough that most blocks end up owning an empty column range.
             (64, 2048),
             (64, 8192),
-            # gfx942 auto-selects r4 from 256*8192 == 2**21 up (pinned in
-            # _GFX942_AUTO_VARIANTS). 257 is odd, so the launcher pads it.
+            # gfx942 auto-selects r4 from 256*8192 == 2**21 up. 257 leaves a
+            # short last tile.
             (256, 8192),
             (257, 8192),
             # s_kv < s_q. A causal mask then puts cu_ends below zero on the