[FlyDSL] [gfx950] Add explicit page strides and 64-bit rebase for FP4 MQA logits

aiter/ops/flydsl/kernels/fused_compress_attn.py

@@ -1153,19 +1153,17 @@ def _phase2_issue_loads(k_i32):
                             fx.Int8,
                         )
                         if const_expr(preshuffle):
-                            # scale [NB, k_tiles, 4, kvbs] u8, with the slot axis
-                            # INTERLEAVED so the mqa-logits reader's packed-dword
-                            # load (4 nt-bytes adjacent) is contiguous:
-                            #   sflat = (slot % 16) * KVS_NTPW + (slot // 16)
-                            # (KVS_NTPW == 4). Matches the op-test reference
-                            # writer `indexer_k_fp4_paged_preshuffle` and the
-                            # packed N_PHYS==1 readers in pa_mqa_logits_fp4*.
+                            # scale [NB, k_tiles, 4, kvbs] u8, with four
+                            # equally-sized token groups interleaved so the
+                            # mqa-logits reader can assemble one packed scale
+                            # dword for each 64-token warp tile.
                             sg_u = fx.Uint32(scale_group_idx)
                             k_tile_s = fx.Int32(sg_u // fx.Uint32(4))
                             group4_s = fx.Int32(sg_u % fx.Uint32(4))
                             slot_u = fx.Uint32(slot_in_block)
-                            sflat = fx.Int32(slot_u % fx.Uint32(16)) * 4 + fx.Int32(
-                                slot_u // fx.Uint32(16)
+                            scale_group_size = fx.Uint32(KVBS // 4)
+                            sflat = fx.Int32(slot_u % scale_group_size) * 4 + fx.Int32(
+                                slot_u // scale_group_size
                             )
                             cs_off = k_tile_s * (4 * KVBS) + group4_s * KVBS + sflat
                         else:

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

@@ -30,6 +30,52 @@
 DEFAULT_BLOCK_THREADS = DEFAULT_NUM_WARPS * WARP_SIZE  # 256
 
 
+def default_varctx_parallel_unit_num(
+    batch_size: int,
+    max_seq_len: int,
+    next_n: int = 1,
+) -> int:
+    """Return the tuned gfx950 persistent-grid size without reading the GPU.
+
+    Context lengths stay device-side so graph replay never synchronizes with
+    the host. The grid can therefore depend only on the host-known row count
+    and maximum sequence length. These buckets are process-level medians from
+    8K/32K/64K/128K sweeps on a 256-CU MI355X; callers can still pass an
+    explicit ``parallel_unit_num`` to override them.
+    """
+    if batch_size < 1 or max_seq_len < 1 or next_n < 1:
+        raise ValueError(
+            "batch_size, max_seq_len, and next_n must all be positive; "
+            f"got {batch_size=}, {max_seq_len=}, {next_n=}."
+        )
+
+    rows = batch_size * next_n
+    if rows <= 1:
+        row_bucket = 0
+    elif rows <= 2:
+        row_bucket = 1
+    elif rows <= 4:
+        row_bucket = 2
+    elif rows <= 8:
+        row_bucket = 3
+    elif rows <= 32:
+        row_bucket = 4
+    else:
+        row_bucket = 5
+
+    if max_seq_len <= 16384:
+        targets = (64, 96, 256, 256, 512, 512)
+    elif max_seq_len <= 49152:
+        targets = (128, 256, 256, 512, 1024, 2048)
+    elif max_seq_len <= 98304:
+        targets = (256, 512, 1024, 1280, 1536, 2048)
+    else:
+        targets = (256, 512, 1280, 1280, 2048, 4096)
+
+    target = max(rows, targets[row_bucket])
+    return ((target + next_n - 1) // next_n) * next_n
+
+
 @triton.jit
 def _varctx_cta_info_kernel(
     ctx_ptr,  # [B] int32
@@ -113,8 +159,9 @@ def compute_varctx_schedule(
 ):
     B = context_lens.shape[0]
     if parallel_unit_num is None:
-        chunks_per_seq = max(1, (max_seq_len + block_k - 1) // block_k)
-        parallel_unit_num = B * next_n * chunks_per_seq
+        parallel_unit_num = default_varctx_parallel_unit_num(
+            B, max_seq_len, next_n=next_n
+        )
     P = parallel_unit_num
     if P % next_n != 0:
         raise ValueError(f"parallel_unit_num={P} must be a multiple of next_n={next_n}")
@@ -134,6 +181,11 @@ def compute_varctx_schedule(
     if cta_info_out is None:
         cta_info = torch.empty(P, 4, dtype=torch.int32, device=dev)
     else:
+        if cta_info_out.shape[0] < P or cta_info_out.shape[1:] != (4,):
+            raise ValueError(
+                f"cta_info_out must have shape (at least {P}, 4), "
+                f"got {tuple(cta_info_out.shape)}."
+            )
         cta_info = cta_info_out
     safe_out = torch.empty(1, dtype=torch.int32, device=dev)
     BLOCK_B = triton.next_power_of_2(max(int(B), 1))
@@ -164,6 +216,9 @@ def build_pa_mqa_logits_fp4_module(
     next_n=1,
     heads=DEFAULT_HEADS,
     head_dim=DEFAULT_HEAD_DIM,
+    kv_page_stride: int | None = None,
+    kv_scale_page_stride: int | None = None,
+    block_table_stride: int | None = None,
 ):
     block_threads_k = num_warps * WARP_SIZE
     head_dim_packed = head_dim // 2
@@ -193,14 +248,22 @@ def build_pa_mqa_logits_fp4_module(
     _stride_q_next_n = heads * head_dim_packed  # bytes per next_n slice
     _stride_q_batch = next_n * _stride_q_next_n  # bytes per batch
     _stride_w_batch = heads
-    _stride_bt = max_blocks_per_seq
+    _stride_bt = (
+        max_blocks_per_seq if block_table_stride is None else block_table_stride
+    )
 
     _kv_chunk_bytes = 16
     _stride_kv_ktile = 4 * kv_block_size * _kv_chunk_bytes  # bytes per K_TILE block
-    _stride_kv_block = k_tiles * _stride_kv_ktile  # bytes per phys block
+    _stride_kv_block = (
+        k_tiles * _stride_kv_ktile if kv_page_stride is None else kv_page_stride
+    )
     # KV_scale: [block_id, K_TILES, K_chunks=4, block_size]
     _stride_kvs_ktile = 4 * kv_block_size  # bytes per K_TILE block
-    _stride_kvs_block = k_tiles * _stride_kvs_ktile
+    _stride_kvs_block = (
+        k_tiles * _stride_kvs_ktile
+        if kv_scale_page_stride is None
+        else kv_scale_page_stride
+    )
 
     QS_DW = (m_tiles + 3) // 4
     qs_pad = QS_DW * 4
@@ -253,8 +316,6 @@ def pa_mqa_logits_fp4_kernel(
         cta_info_bt = fx.rocdl.make_buffer_tensor(cta_info_ptr)
         cta_info_vec = fx.Vector(_load_vec4_i32(cta_info_bt, pid * fx.Int32(4)))
 
-        kv_bt = _i32_buffer(kv_cache_ptr, width=4)
-        kvs_bt = _i32_buffer(kv_scale_ptr, width=1)
         bt_bt = _i32_buffer(kv_indices_ptr, width=1)
 
         ZERO_F = fx.Float32(0.0)
@@ -384,30 +445,51 @@ def _load_phys(c_i32_arg):
                 + lane_mod_16
             )
             bi_base = token_global_base // kv_block_size
-            phys_vec = bt_bt[pid_b * _stride_bt + bi_base]
+            valid = (token_global_base < context_len) & (bi_base < max_blocks_per_seq)
+            safe_bi = valid.select(bi_base, fx.Int32(0))
+            phys_vec = bt_bt[pid_b * _stride_bt + safe_bi]
+            phys_vec = valid.select(phys_vec, fx.Int32(0))
             return _phys_to_list(phys_vec)
 
         def _prefetch_chunk(c_i32_arg, phys_list):
             assert N_TILES_PER_WARP == 4, "packed kvs assumes NTPW=4"
-            assert N_PHYS == 1, "packed kvs assumes N_PHYS=1 (NTPW nts share one phys)"
+            assert N_PHYS == 1, "packed kvs assumes one physical page per warp"
 
             kv_list = []
             kvs_packed_list = []
 
             # ---- KVS packed load: 1 dword per k_tile covering 4 nts ----
-            # Address: phys * stride + k_tile_stride + D*kv_block_size + T*NTPW
-            # (T*NTPW because the host puts 4 nt-bytes adjacent per token-group)
-            phys_shared = phys_list[0]
+            # Rebase the descriptor in i64: BLHNC physical page offsets can
+            # exceed the 4 GiB range of a buffer instruction's byte offset.
+            phys_shared = fx.Int64(_uniform(fx.Int32(phys_list[0])))
+            kv_bt = _i32_buffer(
+                kv_cache_ptr,
+                width=4,
+                byte_offset=phys_shared * fx.Int64(_stride_kv_block),
+            )
+            kvs_bt = _i32_buffer(
+                kv_scale_ptr,
+                width=1,
+                byte_offset=phys_shared * fx.Int64(_stride_kvs_block),
+            )
+            # Scales are interleaved in four equal token groups. A 64-row page
+            # puts this warp's four nt scales in one dword; a 128-row page puts
+            # them across two dwords, so extract and repack the bytes.
+            scale_group = kv_block_size // 4
+            warp_token_base = (warp_id * fx.Int32(64)) % kv_block_size
             for k_tile in range_constexpr(k_tiles):
-                kvs_packed_off_bytes = (
-                    phys_shared * _stride_kvs_block
-                    + fx.Int32(k_tile * _stride_kvs_ktile)
-                    + lane_div_16 * kv_block_size
-                    + lane_mod_16 * fx.Int32(N_TILES_PER_WARP)
-                )
-                # 1 dword/thread; byte offset -> i32 element offset (÷4).
-                kvs_packed = kvs_bt[kvs_packed_off_bytes // 4]
-                kvs_packed_list.append(kvs_packed)
+                packed = fx.Int32(0)
+                for nt in range_constexpr(N_TILES_PER_WARP):
+                    token = warp_token_base + fx.Int32(nt * MFMA_N) + lane_mod_16
+                    word_off_bytes = (
+                        fx.Int32(k_tile * _stride_kvs_ktile)
+                        + lane_div_16 * kv_block_size
+                        + (token % scale_group) * fx.Int32(4)
+                    )
+                    word = kvs_bt[word_off_bytes // 4]
+                    scale = (fx.Int32(word) >> ((token // scale_group) * 8)) & 0xFF
+                    packed = packed | (scale << fx.Int32(nt * 8))
+                kvs_packed_list.append(packed)
 
             # ---- KV loads (unchanged): 1 dwordx4 per (nt, k_tile) ----
             for nt in range_constexpr(N_TILES_PER_WARP):
@@ -417,14 +499,12 @@ def _prefetch_chunk(c_i32_arg, phys_list):
                     + ni_c * fx.Int32(MFMA_N)
                     + lane_mod_16
                 )
-                # No address clamping — OOB tokens read garbage that is later
-                # overwritten by NEG_INF via in_bounds.select on the store path.
+                # Tail pages use physical page zero; windowed stores discard
+                # tokens outside the visible context.
                 token_in_block_c = token_global_c % kv_block_size
-                phys_block_c = phys_list[nt]
                 for k_tile in range_constexpr(k_tiles):
                     kv_off_bytes_c = (
-                        phys_block_c * _stride_kv_block
-                        + fx.Int32(k_tile * _stride_kv_ktile)
+                        fx.Int32(k_tile * _stride_kv_ktile)
                         + lane_div_16 * kv_block_size * _kv_chunk_bytes
                         + token_in_block_c * _kv_chunk_bytes
                     )
@@ -634,6 +714,9 @@ def compile_pa_mqa_logits_fp4(
     next_n: int = 1,
     heads: int = DEFAULT_HEADS,
     head_dim: int = DEFAULT_HEAD_DIM,
+    kv_page_stride: int | None = None,
+    kv_scale_page_stride: int | None = None,
+    block_table_stride: int | None = None,
 ):
     kfn, block_threads = build_pa_mqa_logits_fp4_module(
         block_k=block_k,
@@ -643,6 +726,9 @@ def compile_pa_mqa_logits_fp4(
         next_n=next_n,
         heads=heads,
         head_dim=head_dim,
+        kv_page_stride=kv_page_stride,
+        kv_scale_page_stride=kv_scale_page_stride,
+        block_table_stride=block_table_stride,
     )
 
     @flyc.jit
@@ -692,14 +778,18 @@ def flydsl_pa_mqa_logits_fp4(
     """Decode/varctx FP4 paged MQA logits (gfx950).
 
     ``parallel_unit_num`` is the persistent-grid CTA count; when ``None`` it is
-    auto-derived (cudagraph-safe, no device→host sync) as
-    ``batch * next_n * ceil(max_seq_len / block_k)``, which is a multiple of
-    ``next_n`` and ``>= batch*next_n`` by construction. Pass a smaller explicit
-    value to trade parallelism for fewer no-op CTAs.
+    selected by :func:`default_varctx_parallel_unit_num` from host-known shape
+    buckets (cudagraph-safe, no device→host sync). Pass an explicit value to
+    override the tuned default.
     """
     batch_size, q_next_n, heads, head_dim_packed = q_fp4.shape
     head_dim = head_dim_packed * 2
     max_blocks_per_seq = block_tables.shape[1]
+    if kv_block_size not in (64, 128) or block_k != 256:
+        raise ValueError(
+            "FP4 packed-scale kernels require kv_block_size=64 or 128 and "
+            "block_k=256"
+        )
     if q_next_n != next_n:
         raise ValueError(f"q_fp4 next_n dim ({q_next_n}) != next_n arg ({next_n}).")
 
@@ -729,6 +819,9 @@ def flydsl_pa_mqa_logits_fp4(
         next_n=next_n,
         heads=heads,
         head_dim=head_dim,
+        kv_page_stride=kv_cache.stride(0),
+        kv_scale_page_stride=kv_scale.stride(0),
+        block_table_stride=block_tables.stride(0),
     )
 
     if stream is None:

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

@@ -15,14 +15,18 @@
 _NON_WRITER_LANE_OFF = 1 << 30
 
 
-def _i32_buffer(ptr, width=1):
+def _i32_buffer(ptr, width=1, byte_offset=None):
     """OOB-checked global i32 buffer-tensor over ``ptr`` (mirrors a max_size V#).
 
     ``width`` shapes it ``(N, width)`` so a per-``width`` row can be sliced and
     vector-copied; ``width=1`` gives a flat tensor for scalar ``[idx]`` loads.
+    ``byte_offset`` rebases the global pointer before constructing the resource;
+    buffer instruction offsets alone cannot address beyond 4 GiB.
     """
     src = fx.get_iter(ptr)
     it = fx.recast_iter(fx.PointerType.get(T.i32, src.memspace, 4), src)
+    if byte_offset is not None:
+        it = fx.add_offset(it, byte_offset // fx.Int64(4))
     if width == 1:
         lay = fx.make_layout((1 << 30,), (1,))
     else: