From d6807379f29136220fa8a431df8f6709cf2f2919 Mon Sep 17 00:00:00 2001 From: LiuYinfeng01 Date: Tue, 15 Sep 2026 05:39:22 +0000 Subject: [PATCH 1/5] [FlyDSL][gfx950] Add explicit page strides and 64-bit rebase for FP4 MQA logits Teach pa_mqa_logits_fp4 decode/prefill to honor caller-provided K/scale/block-table page strides and rebase each physical cache page in 64-bit address space before building buffer descriptors. Fixes wrong logits when shared KV pools place pages beyond the 4 GiB range of buffer-instruction byte offsets (reproduced at 4.6 GiB in DeepSeek V4 BLHNC layouts). Assisted-by: OpenAI API coding assistant Signed-off-by: LiuYinfeng01 --- .../kernels/mqa_logits/pa_mqa_logits_fp4.py | 86 +++++++++++++------ .../mqa_logits/pa_mqa_logits_fp4_common.py | 6 +- .../mqa_logits/pa_mqa_logits_fp4_prefill.py | 85 ++++++++++++------ 3 files changed, 122 insertions(+), 55 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py index a2972fe6da5..552b897d5c8 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py @@ -164,43 +164,54 @@ 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 m_tiles = heads // MFMA_M k_tiles = head_dim // 128 # outer K-loop iters (MFMA K=128) - assert ( - head_dim % 128 == 0 - ), f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" + assert head_dim % 128 == 0, ( + f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" + ) assert heads % MFMA_M == 0, f"heads must be a multiple of {MFMA_M}, got {heads}" N_TILES = block_k // MFMA_N - assert ( - N_TILES % num_warps == 0 - ), f"block_k={block_k} → N_TILES={N_TILES} must be multiple of num_warps={num_warps}" + assert N_TILES % num_warps == 0, ( + f"block_k={block_k} → N_TILES={N_TILES} must be multiple of num_warps={num_warps}" + ) N_TILES_PER_WARP = N_TILES // num_warps assert kv_block_size % MFMA_N == 0, ( f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}; " f"sub-tile pages would require splitting one MFMA over multiple page lookups" ) - assert ( - block_k % kv_block_size == 0 - ), f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" + assert block_k % kv_block_size == 0, ( + f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" + ) TILES_PER_BLOCK = kv_block_size // MFMA_N N_PHYS = (N_TILES_PER_WARP + TILES_PER_BLOCK - 1) // TILES_PER_BLOCK _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 +264,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,7 +393,10 @@ 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): @@ -395,13 +407,20 @@ def _prefetch_chunk(c_i32_arg, phys_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), + ) for k_tile in range_constexpr(k_tiles): kvs_packed_off_bytes = ( - phys_shared * _stride_kvs_block - + fx.Int32(k_tile * _stride_kvs_ktile) + fx.Int32(k_tile * _stride_kvs_ktile) + lane_div_16 * kv_block_size + lane_mod_16 * fx.Int32(N_TILES_PER_WARP) ) @@ -417,14 +436,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 ) @@ -502,9 +519,9 @@ def _bperm_xor_add(val, sh): def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): """Process chunk c using prefetched (kv, kvs_packed).""" - assert ( - N_TILES_PER_WARP == 4 - ), "pipelined-nt structure currently hardcoded for NTPW=4" + assert N_TILES_PER_WARP == 4, ( + "pipelined-nt structure currently hardcoded for NTPW=4" + ) # Pre-extract all NTPW nt scales (frees packed register early). kvs_scales = _extract_kvs_scales(kvs_packed_list_in) @@ -634,6 +651,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 +663,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 @@ -700,6 +723,10 @@ def flydsl_pa_mqa_logits_fp4( 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 != 64 or block_k != 256: + raise ValueError( + "FP4 packed-scale kernels require kv_block_size=64, 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 +756,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: diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py index df9d1bef50c..498491ba26c 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_common.py +++ b/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: diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py index de938996939..551a93b35a6 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py @@ -364,44 +364,55 @@ def build_pa_mqa_logits_fp4_prefill_module( num_warps=DEFAULT_NUM_WARPS, 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, ): """Build the ragged-prefill FP4 MQA logits kernel.""" block_threads_k = num_warps * WARP_SIZE m_tiles = heads // MFMA_M k_tiles = head_dim // 128 # outer K-loop iters (MFMA K=128) - assert ( - head_dim % 128 == 0 - ), f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" + assert head_dim % 128 == 0, ( + f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" + ) assert heads % MFMA_M == 0, f"heads must be a multiple of {MFMA_M}, got {heads}" N_TILES = block_k // MFMA_N - assert ( - N_TILES % num_warps == 0 - ), f"block_k={block_k} -> N_TILES={N_TILES} must be multiple of num_warps={num_warps}" + assert N_TILES % num_warps == 0, ( + f"block_k={block_k} -> N_TILES={N_TILES} must be multiple of num_warps={num_warps}" + ) N_TILES_PER_WARP = N_TILES // num_warps - assert ( - kv_block_size % MFMA_N == 0 - ), f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}" - assert ( - block_k % kv_block_size == 0 - ), f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" + assert kv_block_size % MFMA_N == 0, ( + f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}" + ) + assert block_k % kv_block_size == 0, ( + f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" + ) TILES_PER_BLOCK = kv_block_size // MFMA_N N_PHYS = (N_TILES_PER_WARP + TILES_PER_BLOCK - 1) // TILES_PER_BLOCK # block_tables row stride (i32 elements). - _stride_bt = max_blocks_per_seq + _stride_bt = ( + max_blocks_per_seq if block_table_stride is None else block_table_stride + ) # KV preshuffle layout: [block_id, K_TILES, K_chunk=4, kv_block_size, 16] uint8. _kv_chunk_bytes = 16 _stride_kv_ktile = 4 * kv_block_size * _kv_chunk_bytes - _stride_kv_block = k_tiles * _stride_kv_ktile + _stride_kv_block = ( + k_tiles * _stride_kv_ktile if kv_page_stride is None else kv_page_stride + ) # byte stride between consecutive nt tiles inside one kv block (one MFMA_N # row of tokens); used as the per-nt constant `soffset` immediate delta. _stride_kv_ntile = MFMA_N * _kv_chunk_bytes # KV_scale: [block_id, K_TILES, K_chunks=4, kv_block_size] _stride_kvs_ktile = 4 * kv_block_size - _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 + ) _kb_is_pow2 = kv_block_size & (kv_block_size - 1) == 0 _kb_log2 = kv_block_size.bit_length() - 1 @@ -483,8 +494,6 @@ def _uniform(v): local_start = _uniform(cta_info_bt[(fx.Int32(1), fx.Int32(0))]) local_end = _uniform(cta_info_bt[(fx.Int32(1), fx.Int32(1))]) - 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) @@ -613,7 +622,10 @@ def _load_phys(c_i32_arg): + lane_mod_16 ) bi_base = _floordiv_kb(token_local_base) - phys_vec = bt_bt[batch_id * _stride_bt + bi_base] + valid = (token_local_base < local_end) & (bi_base < max_blocks_per_seq) + safe_bi = valid.select(bi_base, fx.Int32(0)) + phys_vec = bt_bt[batch_id * _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): @@ -623,10 +635,19 @@ def _prefetch_chunk(c_i32_arg, phys_list): kv_list = [] kvs_packed_list = [] - phys_shared = phys_list[0] + # BLHNC cache pages can be separated by all layers' storage. Rebase + # the descriptor in i64, leaving only within-page buffer offsets. + 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), + ) kvs_base_off_elems = ( - phys_shared * _stride_kvs_block - + lane_div_16 * kv_block_size + lane_div_16 * kv_block_size + lane_mod_16 * fx.Int32(N_TILES_PER_WARP) ) >> fx.Int32(2) for k_tile in range_constexpr(k_tiles): @@ -643,8 +664,7 @@ def _prefetch_chunk(c_i32_arg, phys_list): ) token_in_block0 = _mod_kb(token_local0) kv_base_off_elems = ( - phys_shared * _stride_kv_block - + lane_div_16 * kv_block_size * _kv_chunk_bytes + lane_div_16 * kv_block_size * _kv_chunk_bytes + token_in_block0 * _kv_chunk_bytes ) >> fx.Int32(2) for nt in range_constexpr(N_TILES_PER_WARP): @@ -722,9 +742,9 @@ def _bperm_xor_add(val, sh): ) def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): - assert ( - N_TILES_PER_WARP == 4 - ), "pipelined-nt structure currently hardcoded for NTPW=4" + assert N_TILES_PER_WARP == 4, ( + "pipelined-nt structure currently hardcoded for NTPW=4" + ) accs_nt0 = ( _issue_nt_mfmas(kv_list_in, kvs_packed_list_in, 0) @@ -835,6 +855,9 @@ def compile_pa_mqa_logits_fp4_prefill( num_warps: int = DEFAULT_NUM_WARPS, 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_prefill_module( block_k=block_k, @@ -843,6 +866,9 @@ def compile_pa_mqa_logits_fp4_prefill( num_warps=num_warps, 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 @@ -894,6 +920,10 @@ def flydsl_pa_mqa_logits_fp4_prefill( total_tokens, 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 != 64 or block_k != 256: + raise ValueError( + "FP4 packed-scale kernels require kv_block_size=64, block_k=256" + ) if (cta_info is None) != (n_ctas is None): raise ValueError("Pass both cta_info and n_ctas, or neither.") @@ -925,6 +955,9 @@ def flydsl_pa_mqa_logits_fp4_prefill( num_warps=num_warps, 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: From 0030d058ca8d959051030631a1691b12b9e57d5b Mon Sep 17 00:00:00 2001 From: LiuYinfeng01 Date: Tue, 15 Sep 2026 11:20:16 +0000 Subject: [PATCH 2/5] test: validate FP4 shared-pool addressing and performance Signed-off-by: LiuYinfeng01 --- .../kernels/mqa_logits/pa_mqa_logits_fp4.py | 30 +- .../mqa_logits/pa_mqa_logits_fp4_prefill.py | 39 +- .../test_flydsl_pa_mqa_logits_fp4_prefill.py | 116 +++- ..._flydsl_pa_mqa_logits_fp4_strided_pages.py | 568 ++++++++++++++++++ 4 files changed, 716 insertions(+), 37 deletions(-) create mode 100644 op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py index 552b897d5c8..1d0877a9347 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py @@ -172,24 +172,24 @@ def build_pa_mqa_logits_fp4_module( head_dim_packed = head_dim // 2 m_tiles = heads // MFMA_M k_tiles = head_dim // 128 # outer K-loop iters (MFMA K=128) - assert head_dim % 128 == 0, ( - f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" - ) + assert ( + head_dim % 128 == 0 + ), f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" assert heads % MFMA_M == 0, f"heads must be a multiple of {MFMA_M}, got {heads}" N_TILES = block_k // MFMA_N - assert N_TILES % num_warps == 0, ( - f"block_k={block_k} → N_TILES={N_TILES} must be multiple of num_warps={num_warps}" - ) + assert ( + N_TILES % num_warps == 0 + ), f"block_k={block_k} → N_TILES={N_TILES} must be multiple of num_warps={num_warps}" N_TILES_PER_WARP = N_TILES // num_warps assert kv_block_size % MFMA_N == 0, ( f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}; " f"sub-tile pages would require splitting one MFMA over multiple page lookups" ) - assert block_k % kv_block_size == 0, ( - f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" - ) + assert ( + block_k % kv_block_size == 0 + ), f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" TILES_PER_BLOCK = kv_block_size // MFMA_N N_PHYS = (N_TILES_PER_WARP + TILES_PER_BLOCK - 1) // TILES_PER_BLOCK @@ -411,11 +411,13 @@ def _prefetch_chunk(c_i32_arg, phys_list): # 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, + kv_cache_ptr, + width=4, byte_offset=phys_shared * fx.Int64(_stride_kv_block), ) kvs_bt = _i32_buffer( - kv_scale_ptr, width=1, + kv_scale_ptr, + width=1, byte_offset=phys_shared * fx.Int64(_stride_kvs_block), ) for k_tile in range_constexpr(k_tiles): @@ -519,9 +521,9 @@ def _bperm_xor_add(val, sh): def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): """Process chunk c using prefetched (kv, kvs_packed).""" - assert N_TILES_PER_WARP == 4, ( - "pipelined-nt structure currently hardcoded for NTPW=4" - ) + assert ( + N_TILES_PER_WARP == 4 + ), "pipelined-nt structure currently hardcoded for NTPW=4" # Pre-extract all NTPW nt scales (frees packed register early). kvs_scales = _extract_kvs_scales(kvs_packed_list_in) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py index 551a93b35a6..37a2330d399 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py @@ -372,23 +372,23 @@ def build_pa_mqa_logits_fp4_prefill_module( block_threads_k = num_warps * WARP_SIZE m_tiles = heads // MFMA_M k_tiles = head_dim // 128 # outer K-loop iters (MFMA K=128) - assert head_dim % 128 == 0, ( - f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" - ) + assert ( + head_dim % 128 == 0 + ), f"head_dim must be a multiple of 128 (MFMA K), got {head_dim}" assert heads % MFMA_M == 0, f"heads must be a multiple of {MFMA_M}, got {heads}" N_TILES = block_k // MFMA_N - assert N_TILES % num_warps == 0, ( - f"block_k={block_k} -> N_TILES={N_TILES} must be multiple of num_warps={num_warps}" - ) + assert ( + N_TILES % num_warps == 0 + ), f"block_k={block_k} -> N_TILES={N_TILES} must be multiple of num_warps={num_warps}" N_TILES_PER_WARP = N_TILES // num_warps - assert kv_block_size % MFMA_N == 0, ( - f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}" - ) - assert block_k % kv_block_size == 0, ( - f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" - ) + assert ( + kv_block_size % MFMA_N == 0 + ), f"kv_block_size={kv_block_size} must be a multiple of MFMA_N={MFMA_N}" + assert ( + block_k % kv_block_size == 0 + ), f"block_k={block_k} must be a multiple of kv_block_size={kv_block_size}" TILES_PER_BLOCK = kv_block_size // MFMA_N N_PHYS = (N_TILES_PER_WARP + TILES_PER_BLOCK - 1) // TILES_PER_BLOCK @@ -639,16 +639,17 @@ def _prefetch_chunk(c_i32_arg, phys_list): # the descriptor in i64, leaving only within-page buffer offsets. phys_shared = fx.Int64(_uniform(fx.Int32(phys_list[0]))) kv_bt = _i32_buffer( - kv_cache_ptr, width=4, + kv_cache_ptr, + width=4, byte_offset=phys_shared * fx.Int64(_stride_kv_block), ) kvs_bt = _i32_buffer( - kv_scale_ptr, width=1, + kv_scale_ptr, + width=1, byte_offset=phys_shared * fx.Int64(_stride_kvs_block), ) kvs_base_off_elems = ( - lane_div_16 * kv_block_size - + lane_mod_16 * fx.Int32(N_TILES_PER_WARP) + lane_div_16 * kv_block_size + lane_mod_16 * fx.Int32(N_TILES_PER_WARP) ) >> fx.Int32(2) for k_tile in range_constexpr(k_tiles): kvs_packed = kvs_bt[ @@ -742,9 +743,9 @@ def _bperm_xor_add(val, sh): ) def _compute_chunk(kv_list_in, kvs_packed_list_in, c_i32_arg, nt0_accs_in=None): - assert N_TILES_PER_WARP == 4, ( - "pipelined-nt structure currently hardcoded for NTPW=4" - ) + assert ( + N_TILES_PER_WARP == 4 + ), "pipelined-nt structure currently hardcoded for NTPW=4" accs_nt0 = ( _issue_nt_mfmas(kv_list_in, kvs_packed_list_in, 0) diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py index 9e941654800..b32d1194d61 100755 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py @@ -30,6 +30,7 @@ """ import argparse +import statistics import torch @@ -314,10 +315,10 @@ def run_fp4(): assert oob_ok, "OOB cells were not left at -inf" if not bench: - return + return None _, us_fp4 = run_perftest(run_fp4, num_iters=iters, num_warmup=warmup) - _bench_vs_atom( + return _bench_vs_atom( kv_bf16, slot_mapping, block_tables, @@ -333,7 +334,9 @@ def run_fp4(): kv_block_size, t_max, bs, + out, us_fp4, + cos_bf16, iters, warmup, ) @@ -355,7 +358,9 @@ def _bench_vs_atom( kv_block_size, t_max, bs, + fp4_out, us_fp4, + cos_fp4_bf16, iters, warmup, ): @@ -426,12 +431,15 @@ def atom_logits(): atom_out = atom_logits() torch.cuda.synchronize() - vo, vr = [], [] + vo, vr, vf = [], [], [] for r in range(total_tokens): b, e = int(row_to_batch[r]), int(le[r]) vo.append(atom_out[r, cu[b] : cu[b] + e]) vr.append(ref_bf16[r, :e]) - cos_atom = _cos(torch.cat(vo), torch.cat(vr)).item() + vf.append(fp4_out[r, :e]) + atom_values = torch.cat(vo) + cos_atom = _cos(atom_values, torch.cat(vr)).item() + cos_fp4_fp8 = _cos(torch.cat(vf), atom_values).item() _, us_gather = run_perftest( cp_gather_indexer_k_quant_cache, @@ -455,8 +463,70 @@ def atom_logits(): if atom_ok: print(" {:<28} | {:>10.2f}".format("ATOM cp_gather", us_gather)) print(" {:<28} | {:>10.2f}".format("ATOM cp_gather+fp8_logits", us_atom)) + print(f" [accuracy] FP4 vs ATOM fp8 cos={cos_fp4_fp8:.6f}") print(f" [accuracy] ATOM fp8 vs bf16 ref cos={cos_atom:.6f}") print(f" speedup (ATOM total / FP4) = {us_atom / us_fp4:.2f}x") + if not atom_ok: + return None + return { + "fp4_us": us_fp4, + "fp8_gather_us": us_gather, + "fp8_total_us": us_atom, + "speedup": us_atom / us_fp4, + "fp4_bf16_cos": cos_fp4_bf16, + "fp8_bf16_cos": cos_atom, + "fp4_fp8_cos": cos_fp4_fp8, + } + + +def _print_standard_bench_summary(samples): + print("\n" + "=" * 126) + print("Standard prefill indexer sweep (FP8 gather+logits vs FP4 direct-paged)") + print("=" * 126) + print( + " {:<11} | {:>3} | {:>5} | {:>6} | {:>8} | {:>8} | {:>8} | " + "{:>7} | {:>9} | {:>9} | {:>9} | {:>7}".format( + "shape", + "bs", + "n_q", + "ctx", + "fp4_us", + "gather", + "fp8_us", + "speedup", + "fp4_cos", + "fp8_cos", + "fp4/fp8", + "repeats", + ) + ) + print(" " + "-" * 122) + for shape, rows in samples: + row = { + key: statistics.median(sample[key] for sample in rows) + for key in ( + "fp4_us", + "fp8_gather_us", + "fp8_total_us", + "speedup", + "fp4_bf16_cos", + "fp8_bf16_cos", + "fp4_fp8_cos", + ) + } + print( + " {name:<11} | {bs:>3} | {n_q:>5} | {ctx:>6} | " + "{fp4_us:>8.2f} | {fp8_gather_us:>8.2f} | {fp8_total_us:>8.2f} | " + "{speedup:>6.2f}x | {fp4_bf16_cos:>9.6f} | {fp8_bf16_cos:>9.6f} | " + "{fp4_fp8_cos:>9.6f} | {repeats:>7}".format( + name=shape[0], + bs=shape[1], + n_q=shape[2], + ctx=shape[3], + repeats=len(rows), + **row, + ) + ) print() @@ -698,6 +768,11 @@ def _print_varqlen_perf_summary(): def main(): ap = argparse.ArgumentParser() ap.add_argument("--bench", action="store_true") + ap.add_argument( + "--bench-sweep", + action="store_true", + help="Run standard 2K/8K/32K prefill indexer shapes", + ) ap.add_argument("--bs", type=int, default=4) ap.add_argument("--ctx", type=int, default=2048) ap.add_argument("--n_q", type=int, default=64) @@ -705,6 +780,12 @@ def main(): ap.add_argument("--head_dim", type=int, default=128) ap.add_argument("--iters", type=int, default=50) ap.add_argument("--warmup", type=int, default=10) + ap.add_argument( + "--bench-repeats", + type=int, + default=3, + help="Independent repetitions per standard sweep shape", + ) args = ap.parse_args() if get_arch() != "gfx950": @@ -734,6 +815,33 @@ def main(): iters=args.iters, warmup=args.warmup, ) + if args.bench_sweep: + standard_shapes = ( + ("small", 1, 32, 2048), + ("medium", 4, 64, 2048), + ("long", 4, 64, 8192), + ("very-long", 8, 16, 32768), + ) + samples = [] + for shape in standard_shapes: + _, bs, n_q, ctx = shape + rows = [] + for repeat in range(args.bench_repeats): + result = run_case( + bs, + [[ctx] * n_q for _ in range(bs)], + heads=args.heads, + head_dim=args.head_dim, + seed=7 + repeat, + bench=True, + iters=args.iters, + warmup=args.warmup, + ) + if result is not None: + rows.append(result) + torch.cuda.empty_cache() + samples.append((shape, rows)) + _print_standard_bench_summary(samples) # ── Variable-qlen (per-batch MTP via cu_seq_q) sweep + TFLOPS/bandwidth ── print("=" * 80) diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py b/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py new file mode 100644 index 00000000000..7d464f655dd --- /dev/null +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py @@ -0,0 +1,568 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FP4 MQA shared-pool addressing regression tests (gfx950). + +The original constant 8192-logit regression at byte offset 4,608,000,000 +(4.608 GB, about 4.292 GiB) is retained. Random tests additionally compare +shared-pool pages to same-byte contiguous pages AND an independent dequantized +FP4 oracle. Physical pages straddle 2 GiB and 4 GiB, with shuffled mappings, +padded block-table rows, nonzero storage offsets, and tail/empty windows. +Pools are allocated sequentially; the largest backing store is about 4.61 GB. + +Usage: + python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py --graph + python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py --case small \ + --graph --sensitivity --benchmark --json-output results.json + +Run the same file with separate before/after AITER source mounts and distinct +--label values. Benchmark numbers time the AITER public API, not vLLM gather. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import platform +import subprocess +from pathlib import Path + +import torch + +from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4 import ( + flydsl_pa_mqa_logits_fp4, +) +from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4_prefill import ( + flydsl_pa_mqa_logits_fp4_prefill, +) +from aiter.ops.triton.utils._triton.arch_info import get_arch +from aiter.test_common import run_perftest + +try: + from test_flydsl_pa_mqa_logits_fp4_prefill import ( + fp4_dequant_e2m1_with_e8m0, + fp4_quant_e2m1_with_e8m0, + indexer_k_fp4_paged_preshuffle, + quant_q_fp4_preshuffle, + ref_prefill_logits, + ) +except ModuleNotFoundError as exc: + if exc.name != "test_flydsl_pa_mqa_logits_fp4_prefill": + raise + from op_tests.test_flydsl_pa_mqa_logits_fp4_prefill import ( + fp4_dequant_e2m1_with_e8m0, + fp4_quant_e2m1_with_e8m0, + indexer_k_fp4_paged_preshuffle, + quant_q_fp4_preshuffle, + ref_prefill_logits, + ) + +dev = "cuda" +HEAD_DIM = 128 +KV_BLOCK_SIZE = 64 +PAGE_BYTES = KV_BLOCK_SIZE * 68 +BLOCK_K = 256 +HEADS = 64 +EXPECTED_LOGIT = 8192.0 +GUARD = 123456.0 +SEED = 5518 +# At each power-of-two boundary, used pages are on both sides of the boundary. +CASES = { + "small": (37, 8192), + "2gib": ((1 << 31) // 8192 - 4, 8192), + "4gib": ((1 << 32) // 8192 - 4, 8192), + "original": (20000, 230400), +} + + +def vllm_fp4_cache_views( + kv_cache: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Split vLLM (num_blocks, 64, 68) layout into AITER preshuffle views.""" + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 + num_blocks, block_size, row_bytes = kv_cache.shape + assert block_size == KV_BLOCK_SIZE and row_bytes == 68 + assert kv_cache.stride(2) == 1 and kv_cache.stride(1) == 68 + page_stride = kv_cache.stride(0) + assert page_stride >= block_size * row_bytes + values = torch.as_strided( + kv_cache, + (num_blocks, 1, 4, block_size, 16), + (page_stride, block_size * 64, block_size * 16, 16, 1), + ) + scales = torch.as_strided( + kv_cache, + (num_blocks, 1, 4, block_size), + (page_stride, block_size * 4, block_size, 1), + storage_offset=kv_cache.storage_offset() + block_size * 64, + ) + return values, scales + + +def padded_table(entries: torch.Tensor) -> torch.Tensor: + """Poison padding and add both a row stride and a nonzero storage offset.""" + rows, columns = entries.shape + storage = torch.full( + (rows + 1, columns + 11), 2147483647, dtype=torch.int32, device=dev + ) + table = storage[1:, 3 : 3 + columns] + table.copy_(entries) + assert table.stride(0) > columns and table.storage_offset() > 0 + return table + + +def allocate_pool(page_base, page_stride, dense_cache): + """Copy only used page bytes, never initialize the multi-GB unused pool.""" + pages = dense_cache.shape[0] + layer_offset = 3 * PAGE_BYTES + 128 + pool_pages = page_base + pages + storage_bytes = layer_offset + (pool_pages - 1) * page_stride + PAGE_BYTES + storage = torch.empty(storage_bytes, dtype=torch.uint8, device=dev) + cache = storage.as_strided( + (pool_pages, KV_BLOCK_SIZE, 68), (page_stride, 68, 1), layer_offset + ) + # Page zero is a valid masked-load target, and a deterministic wrong-address + # target; the correct result never depends on uninitialized page-zero bytes. + cache[0].zero_() + cache[page_base : page_base + pages].copy_(dense_cache) + return cache, storage_bytes + + +def make_inputs(rows, context, constant=False): + """Three requests, distinct per-block scales and signed random FP4 data.""" + generator = torch.Generator(device=dev).manual_seed(SEED) + requests = 3 + # Kernel loads page metadata in four-page chunks, including masked tails. + pages_per_request = (context + BLOCK_K - 1) // BLOCK_K * 4 + tokens = pages_per_request * KV_BLOCK_SIZE + pages = requests * pages_per_request + logical_pages = torch.randperm(pages, generator=generator, device=dev).reshape( + requests, pages_per_request + ) + table = padded_table(logical_pages.to(torch.int32)) + dense_cache = torch.zeros((pages, KV_BLOCK_SIZE, 68), dtype=torch.uint8, device=dev) + dense_values, dense_scales = vllm_fp4_cache_views(dense_cache) + if constant: + q_source = torch.ones((rows + 7, HEADS, HEAD_DIM), device=dev) + k_source = torch.ones((requests, tokens, HEAD_DIM), device=dev) + weights = torch.ones((rows, HEADS), dtype=torch.bfloat16, device=dev) + else: + + def random_scaled(shape): + values = torch.randn(shape, generator=generator, device=dev) + exponents = torch.randint( + -2, + 3, + (*shape[:-1], HEAD_DIM // 32, 1), + generator=generator, + device=dev, + ) + return (values.reshape(*shape[:-1], 4, 32) * (2.0**exponents)).reshape( + shape + ) + + q_source = random_scaled((rows + 7, HEADS, HEAD_DIM)) + k_source = random_scaled((requests, tokens, HEAD_DIM)) + weights = ( + 0.25 + torch.rand((rows, HEADS), generator=generator, device=dev) + ).to(torch.bfloat16) + q_source = q_source.to(torch.bfloat16) + k_source = k_source.to(torch.bfloat16) + q_storage, qs_storage = quant_q_fp4_preshuffle(q_source) + q, qs = q_storage[7:], qs_storage[7:] + _, q_e8 = fp4_quant_e2m1_with_e8m0(q_source[7:]) + q_dq = fp4_dequant_e2m1_with_e8m0(q, q_e8) + k_packed, k_e8 = fp4_quant_e2m1_with_e8m0(k_source) + k_dq = fp4_dequant_e2m1_with_e8m0(k_packed, k_e8) + slots = ( + logical_pages[:, :, None] * KV_BLOCK_SIZE + + torch.arange(KV_BLOCK_SIZE, device=dev) + ).reshape(-1) + indexer_k_fp4_paged_preshuffle( + k_source.reshape(-1, HEAD_DIM), + slots, + dense_values, + dense_scales, + KV_BLOCK_SIZE, + ) + if constant: + # Keep the original literal 0x22 / E8M0=127 byte pattern, not merely + # another quantization of the same real-valued input. + dense_values.fill_(0x22) + dense_scales.fill_(127) + q.fill_(0x22) + qs.fill_(127) + else: + assert q_e8.unique().numel() > 1 and k_e8.unique().numel() > 1 + mapping = torch.arange(rows, device=dev, dtype=torch.int32) % requests + return dense_cache, table, q, qs, weights, q_dq, k_dq, mapping + + +def guarded_output(rows, context): + backing = torch.full((rows + 2, context + 19), GUARD, device=dev) + out = backing[1:-1, 7 : 7 + context] + return backing, out + + +def verify_output(out, backing, reference, *, exact=False): + valid = torch.isfinite(reference) + assert bool(torch.isneginf(out[~valid]).all()), "invalid logits must be -inf" + torch.testing.assert_close( + out[valid], + reference[valid], + rtol=0 if exact else 2e-4, + atol=0 if exact else 2e-3, + ) + assert bool((backing[0] == GUARD).all()) + assert bool((backing[-1] == GUARD).all()) + assert bool((backing[1:-1, :7] == GUARD).all()) + assert bool((backing[1:-1, out.shape[1] + 7 :] == GUARD).all()) + + +def run_case( + name, + *, + graph=False, + sensitivity=False, + benchmark=False, + rows=12, + context=333, + iters=50, + warmup=10, + constant=False, + paths=("prefill", "decode"), + layouts=("dense", "strided"), +): + page_base, page_stride = CASES[name] + dense_cache, dense_table, q, qs, weights, q_dq, k_dq, mapping = make_inputs( + rows, context, constant + ) + dense_values, dense_scales = vllm_fp4_cache_views(dense_cache) + # Stock kernels expect separately contiguous K and scale arrays. Preserve + # every payload byte while removing only the shared-pool storage layout. + dense_values = dense_values.contiguous() + dense_scales = dense_scales.contiguous() + dense_table = dense_table.contiguous() + cache = values = scales = strided_table = None + storage_bytes = 0 + if "strided" in layouts: + cache, storage_bytes = allocate_pool(page_base, page_stride, dense_cache) + strided_table = padded_table(dense_table + page_base) + values, scales = vllm_fp4_cache_views(cache) + assert torch.equal( + values[page_base : page_base + len(dense_cache)], dense_values + ) + assert torch.equal( + scales[page_base : page_base + len(dense_cache)], dense_scales + ) + records = [] + for path in paths: + ends_list = [context, context - 1, min(65, context), 0] + ends = torch.tensor( + [ends_list[r % 4] for r in range(rows)], dtype=torch.int32, device=dev + ) + starts = torch.zeros_like(ends) + if path == "prefill": + starts[1::4] = min(17, context - 1) + starts[2::4] = ends[2::4] # nonzero empty windows + table_dense, table_strided = dense_table, strided_table + else: + table_dense = dense_table[mapping.long()].contiguous() + table_strided = ( + padded_table(strided_table[mapping.long()]) + if strided_table is not None + else None + ) + reference = ref_prefill_logits( + q_dq, k_dq, weights, mapping, starts, ends, context + ) + if constant: + assert bool((reference[torch.isfinite(reference)] == EXPECTED_LOGIT).all()) + outputs = {} + for layout, kv, ks, table in ( + ("dense", dense_values, dense_scales, table_dense), + ("strided", values, scales, table_strided), + ): + if layout not in layouts: + continue + backing, out = guarded_output(rows, context) + + def launch( + path=path, + kv=kv, + ks=ks, + table=table, + starts=starts, + ends=ends, + out=out, + ): + if path == "prefill": + return flydsl_pa_mqa_logits_fp4_prefill( + q, + qs, + kv, + ks, + table, + weights, + mapping, + starts, + ends, + context, + parallel_unit_num=max(512, rows), + out=out, + ) + return flydsl_pa_mqa_logits_fp4( + q[:, None], + qs[:, None], + kv, + ks, + table, + weights, + ends, + context, + kv_block_size=KV_BLOCK_SIZE, + block_k=BLOCK_K, + parallel_unit_num=max(512, rows), + out=out, + ) + + launch() + torch.cuda.synchronize() + verify_output(out, backing, reference, exact=constant) + outputs[layout] = out.clone() + modes = [("eager", launch)] + if graph: + graph_obj = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_obj): + launch() + # Poison the entire logical output between replays, so a no-op + # or stale captured output cannot pass either masks or values. + for _ in range(2): + out.fill_(GUARD) + graph_obj.replay() + torch.cuda.synchronize() + verify_output(out, backing, reference, exact=constant) + modes.append(("graph_replay", graph_obj.replay)) + if sensitivity and not constant and layout == "strided": + valid = torch.isfinite(reference) + saved_table = table.clone() + table.copy_(table.roll(1, dims=1)) + try: + launch() + torch.cuda.synchronize() + assert not torch.allclose( + out[valid], reference[valid], rtol=2e-4, atol=2e-3 + ), "page-map mutation was undetected" + finally: + table.copy_(saved_table) + used_scales = scales[page_base : page_base + len(dense_cache)] + saved_scales = used_scales.clone() + used_scales.copy_(used_scales.roll(1, dims=-1)) + try: + launch() + torch.cuda.synchronize() + assert not torch.allclose( + out[valid], reference[valid], rtol=2e-4, atol=2e-3 + ), "scale-permutation mutation was undetected" + finally: + used_scales.copy_(saved_scales) + launch() + torch.cuda.synchronize() + verify_output(out, backing, reference) + for mode, call in modes: + result = { + "case": name, + "pattern": "constant" if constant else "random", + "path": path, + "layout": layout, + "mode": mode, + "rows": rows, + "requests": 3, + "context": context, + "heads": HEADS, + "head_dim": HEAD_DIM, + "kv_block_size": KV_BLOCK_SIZE, + "block_k": BLOCK_K, + "q_shape": list(q.shape), + "kv_shape": list(kv.shape), + "page_stride_bytes": kv.stride(0), + "scale_page_stride_bytes": ks.stride(0), + "block_table_shape": list(table.shape), + "block_table_stride": table.stride(0), + "block_table_storage_offset": table.storage_offset(), + "q_storage_offset_bytes": q.storage_offset(), + "q_scale_storage_offset_bytes": qs.storage_offset(), + "output_stride": out.stride(0), + "output_storage_offset": out.storage_offset(), + "cache_storage_offset_bytes": ( + cache.storage_offset() if layout == "strided" else 0 + ), + "scale_storage_offset_bytes": ks.storage_offset(), + "physical_page_offset_min_bytes": ( + page_base * page_stride if layout == "strided" else 0 + ), + "physical_page_offset_max_bytes": ( + (page_base + len(dense_cache) - 1) * page_stride + if layout == "strided" + else (len(dense_cache) - 1) * kv.stride(0) + ), + "pool_allocation_bytes": ( + storage_bytes if layout == "strided" else dense_cache.numel() + ), + "correctness": "passed", + "sensitivity_checked": sensitivity + and not constant + and layout == "strided", + } + if benchmark: + _, latency = run_perftest( + call, + num_iters=iters, + num_warmup=warmup, + num_rotate_args=1, + use_cuda_event=True, + ) + result["latency_us"] = float(latency) + verify_output(out, backing, reference, exact=constant) + records.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + # Release graphs before moving to another layout/large allocation. + if graph: + del graph_obj + del modes, call + if len(outputs) == 2: + torch.testing.assert_close( + outputs["strided"], outputs["dense"], rtol=0, atol=0 + ) + return records + + +def test_small_stride_prefill_decode(graph: bool = False) -> None: + run_case("small", graph=graph, sensitivity=True) + + +def test_large_offset_prefill_decode(graph: bool = False) -> None: + """Original 4.608 GB (4.292 GiB) regression plus power-of-two boundaries.""" + for name in ("2gib", "4gib", "original"): + run_case(name, graph=graph) + gc.collect() + torch.cuda.empty_cache() + run_case("original", graph=graph, rows=512, context=2048, constant=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--graph", + action="store_true", + help="Capture and replay BOTH prefill and decode", + ) + parser.add_argument("--case", choices=("all", *CASES), default="all") + parser.add_argument("--path", choices=("both", "prefill", "decode"), default="both") + parser.add_argument( + "--layout", + choices=("both", "dense", "strided"), + default="both", + help="Dense is stock-compatible and avoids shared-pool allocation; both compares identical payload bytes", + ) + parser.add_argument( + "--sensitivity", + action="store_true", + help="Require wrong page-map and scale permutations to disagree", + ) + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--rows", type=int, default=12) + parser.add_argument("--context", type=int, default=333) + parser.add_argument("--iters", type=int, default=50) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--label", default="unlabelled") + parser.add_argument( + "--source-revision", + help="Host-verified SHA when mount git metadata is unavailable", + ) + parser.add_argument( + "--runtime-label", + help="Container image/build provenance supplied by the runner", + ) + parser.add_argument("--json-output", type=Path) + args = parser.parse_args() + if args.rows < 4 or args.context < 65 or args.iters < 1 or args.warmup < 1: + parser.error("rows >= 4, context >= 65, iters >= 1, warmup >= 1 required") + if get_arch() != "gfx950": + raise SystemExit(f"gfx950 required; current architecture: {get_arch()}") + source = Path( + __import__(flydsl_pa_mqa_logits_fp4.__module__, fromlist=["__file__"]).__file__ + ).resolve() + revision = subprocess.run( + ["git", "-C", str(source.parent), "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=False, + ) + report = { + "metadata": { + "label": args.label, + "torch": torch.__version__, + "hip": torch.version.hip, + "python": platform.python_version(), + "gpu": torch.cuda.get_device_name(), + "arch": get_arch(), + "source": str(source), + "revision": args.source_revision or revision.stdout.strip() or None, + "revision_origin": ( + "runner" + if args.source_revision + else ("git" if revision.returncode == 0 else "unavailable") + ), + "runtime_label": args.runtime_label, + "seed": SEED, + "iterations": args.iters, + "warmup": args.warmup, + "timer": "aiter.test_common.run_perftest CUDA events; fixed inputs, no rotation", + "timing_scope": "public API including internal schedule and output reset; graph mode is one captured API replay", + "rtol": 2e-4, + "atol": 2e-3, + }, + "results": [], + } + for name in CASES if args.case == "all" else (args.case,): + report["results"].extend( + run_case( + name, + graph=args.graph, + sensitivity=args.sensitivity, + benchmark=args.benchmark, + rows=args.rows, + context=args.context, + iters=args.iters, + warmup=args.warmup, + paths=("prefill", "decode") if args.path == "both" else (args.path,), + layouts=( + ("dense", "strided") if args.layout == "both" else (args.layout,) + ), + ) + ) + gc.collect() + torch.cuda.empty_cache() + if args.case == "all": + report["results"].extend( + run_case( + "original", + graph=args.graph, + rows=512, + context=2048, + constant=True, + paths=("prefill", "decode") if args.path == "both" else (args.path,), + layouts=( + ("dense", "strided") if args.layout == "both" else (args.layout,) + ), + ) + ) + if args.json_output: + args.json_output.write_text(json.dumps(report, indent=2) + "\n") + print("PASS: all requested FP4 shared-pool comparisons and guards") + + +if __name__ == "__main__": + main() From 063f3eb5dce1a5780a841fe241a9da6d8d6d98fb Mon Sep 17 00:00:00 2001 From: LiuYinfeng01 Date: Thu, 17 Sep 2026 09:19:24 +0000 Subject: [PATCH 3/5] test: make FP4 strided-page validation CI-safe Exit successfully on non-gfx950 runners and enable graph replay plus page-map/scale sensitivity checks in the default CI invocation. Signed-off-by: LiuYinfeng01 --- ..._flydsl_pa_mqa_logits_fp4_strided_pages.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py b/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py index 7d464f655dd..920b75ca152 100644 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py @@ -11,9 +11,11 @@ Pools are allocated sequentially; the largest backing store is about 4.61 GB. Usage: - python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py --graph + python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py --case small \ - --graph --sensitivity --benchmark --json-output results.json + --benchmark --json-output results.json + python op_tests/test_flydsl_pa_mqa_logits_fp4_strided_pages.py --case small \ + --no-graph --no-sensitivity Run the same file with separate before/after AITER source mounts and distinct --label values. Benchmark numbers time the AITER public API, not vLLM gather. @@ -455,8 +457,9 @@ def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--graph", - action="store_true", - help="Capture and replay BOTH prefill and decode", + action=argparse.BooleanOptionalAction, + default=True, + help="Capture and replay BOTH prefill and decode (enabled by default)", ) parser.add_argument("--case", choices=("all", *CASES), default="all") parser.add_argument("--path", choices=("both", "prefill", "decode"), default="both") @@ -468,8 +471,9 @@ def main() -> None: ) parser.add_argument( "--sensitivity", - action="store_true", - help="Require wrong page-map and scale permutations to disagree", + action=argparse.BooleanOptionalAction, + default=True, + help="Require wrong page-map and scale permutations to disagree (enabled by default)", ) parser.add_argument("--benchmark", action="store_true") parser.add_argument("--rows", type=int, default=12) @@ -489,8 +493,10 @@ def main() -> None: args = parser.parse_args() if args.rows < 4 or args.context < 65 or args.iters < 1 or args.warmup < 1: parser.error("rows >= 4, context >= 65, iters >= 1, warmup >= 1 required") - if get_arch() != "gfx950": - raise SystemExit(f"gfx950 required; current architecture: {get_arch()}") + arch = get_arch() + if arch != "gfx950": + print(f"[skip] this kernel only supports gfx950 (current: {arch}).") + return source = Path( __import__(flydsl_pa_mqa_logits_fp4.__module__, fromlist=["__file__"]).__file__ ).resolve() From 6ab04df300e6b8b6b7401c2a5ab2dcd0ec2b8da6 Mon Sep 17 00:00:00 2001 From: LiuYinfeng01 Date: Thu, 17 Sep 2026 13:28:14 +0000 Subject: [PATCH 4/5] perf: tune FP4 varctx persistent CTA grids Select a graph-stable CTA count from host-known batch and context buckets tuned on MI355X, while preserving explicit caller overrides. Cover the 8K/32K/64K/128K schedule table and MTP divisibility constraints. Signed-off-by: LiuYinfeng01 --- .../kernels/mqa_logits/pa_mqa_logits_fp4.py | 63 +++++++++++++++++-- op_tests/test_flydsl_pa_mqa_logits_fp4.py | 33 ++++++++-- 2 files changed, 85 insertions(+), 11 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py index 1d0877a9347..8ae4bfa0098 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py +++ b/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)) @@ -717,10 +769,9 @@ 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 diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4.py b/op_tests/test_flydsl_pa_mqa_logits_fp4.py index 8e24f851912..6034579af6f 100755 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4.py @@ -14,6 +14,9 @@ import torch from aiter.ops.flydsl import flydsl_pa_mqa_logits_fp4 +from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4 import ( + default_varctx_parallel_unit_num, +) from aiter.ops.triton.utils._triton.arch_info import get_arch from aiter.test_common import checkAllclose, run_perftest @@ -468,9 +471,8 @@ def test_pa_mqa_logits_fp4_qfp4_kvfp4( # The persistent-grid schedule has S = parallel_unit_num // next_n batch # slots; if batch_size exceeds S the surplus batches are silently dropped # (their out stays -inf -> NaN cosine). For an explicit value grow the grid - # so every (batch, next_n) gets at least one slot; when None, let the - # scheduler auto-derive a cudagraph-safe grid (= batch*next_n*ceil(t_max/ - # block_k)), which already satisfies both constraints. + # so every (batch, next_n) gets at least one slot; when None, use the tuned + # cudagraph-safe gfx950 grid, which already satisfies both constraints. if parallel_unit_num is not None: parallel_unit_num = max(parallel_unit_num, batch_size * next_n) safe, cta_info, total_ctas = compute_varctx_schedule( @@ -643,6 +645,27 @@ def _print_perf_summary(): print() +def test_default_varctx_parallel_unit_num(): + expected = { + 8192: (64, 96, 256, 256, 512, 512), + 32768: (128, 256, 256, 512, 1024, 2048), + 65536: (256, 512, 1024, 1280, 1536, 2048), + 131072: (256, 512, 1280, 1280, 2048, 4096), + } + batches = (1, 2, 4, 8, 16, 64) + for context, targets in expected.items(): + actual = tuple( + default_varctx_parallel_unit_num(batch, context) for batch in batches + ) + assert actual == targets, f"{context=}: {actual=} != {targets=}" + + # MTP rows share the same tuning buckets; the launch grid remains a + # multiple of next_n and can never drop a row. + for next_n in (2, 3, 5): + ctas = default_varctx_parallel_unit_num(3, 131072, next_n=next_n) + assert ctas >= 3 * next_n and ctas % next_n == 0 + + def main(): parser = argparse.ArgumentParser( description="MQA Logits (Q FP4, KV FP4) decode Test + Benchmark (gfx950)", @@ -673,8 +696,7 @@ def main(): "--parallel_unit_num", type=int, default=None, - help="target CTA count for host schedule " - "(default: auto = batch*next_n*ceil(max_seq_len/block_k))", + help="target CTA count for host schedule (default: tuned gfx950 auto grid)", ) parser.add_argument( "--next_n", @@ -695,6 +717,7 @@ def main(): help=f"Per-head dim (multiple of 128). Default {DEFAULT_HEAD_DIM}.", ) args = parser.parse_args() + test_default_varctx_parallel_unit_num() if get_arch() != "gfx950": print(f"[skip] this kernel only supports gfx950 (current: {get_arch()}).") From 7db41be280925ae0a9f0961671077aac8e131bb4 Mon Sep 17 00:00:00 2001 From: LiuYinfeng01 Date: Fri, 18 Sep 2026 11:10:36 +0000 Subject: [PATCH 5/5] feat: support 128-token FP4 indexer pages Signed-off-by: LiuYinfeng01 --- .../ops/flydsl/kernels/fused_compress_attn.py | 16 +++---- .../kernels/mqa_logits/pa_mqa_logits_fp4.py | 32 +++++++++----- .../mqa_logits/pa_mqa_logits_fp4_prefill.py | 31 ++++++++----- op_tests/test_flydsl_pa_mqa_logits_fp4.py | 44 ++++++++++++------- .../test_flydsl_pa_mqa_logits_fp4_prefill.py | 10 ++++- 5 files changed, 87 insertions(+), 46 deletions(-) diff --git a/aiter/ops/flydsl/kernels/fused_compress_attn.py b/aiter/ops/flydsl/kernels/fused_compress_attn.py index 2eeeab7d477..32b0821fccc 100644 --- a/aiter/ops/flydsl/kernels/fused_compress_attn.py +++ b/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: diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py index 8ae4bfa0098..42e8fa1e775 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4.py @@ -453,7 +453,7 @@ def _load_phys(c_i32_arg): 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 = [] @@ -472,15 +472,24 @@ def _prefetch_chunk(c_i32_arg, phys_list): 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 = ( - 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): @@ -776,9 +785,10 @@ def flydsl_pa_mqa_logits_fp4( 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 != 64 or block_k != 256: + if kv_block_size not in (64, 128) or block_k != 256: raise ValueError( - "FP4 packed-scale kernels require kv_block_size=64, block_k=256" + "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}).") diff --git a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py index 37a2330d399..4ae2ab6b579 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py @@ -630,7 +630,7 @@ def _load_phys(c_i32_arg): 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 = [] @@ -648,14 +648,24 @@ def _prefetch_chunk(c_i32_arg, phys_list): width=1, byte_offset=phys_shared * fx.Int64(_stride_kvs_block), ) - kvs_base_off_elems = ( - lane_div_16 * kv_block_size + lane_mod_16 * fx.Int32(N_TILES_PER_WARP) - ) >> fx.Int32(2) + # 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 = kvs_bt[ - kvs_base_off_elems + fx.Int32(k_tile * _stride_kvs_ktile // 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) ni0 = warp_id * fx.Int32(N_TILES_PER_WARP) token_local0 = ( @@ -921,9 +931,10 @@ def flydsl_pa_mqa_logits_fp4_prefill( total_tokens, 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 != 64 or block_k != 256: + if kv_block_size not in (64, 128) or block_k != 256: raise ValueError( - "FP4 packed-scale kernels require kv_block_size=64, block_k=256" + "FP4 packed-scale kernels require kv_block_size=64 or 128 and " + "block_k=256" ) if (cta_info is None) != (n_ctas is None): diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4.py b/op_tests/test_flydsl_pa_mqa_logits_fp4.py index 6034579af6f..d57095aade0 100755 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4.py @@ -134,23 +134,22 @@ def create_paged_preshuffle_kv_fp4(kv_bf16, kv_block_size, num_blocks, block_tab .contiguous() .view(batch * t_blocks, k_tiles, 4, kv_block_size, 16) ) - # KVS_NTPW: nt-bytes packed together for the kernel's packed dword load - # (4 ubyte → 1 dword). Per (D=lane_div_16, T=lane_mod_16), bytes for nts - # 0..KVS_NTPW-1 are adjacent so one thread dword-loads all 4 nts. + # Interleave four equally sized token groups so a warp's four scale bytes + # can be reconstructed from one dword per half-page. assert kv_block_size % KVS_NTPW == 0 - kv_e8m0_perm = ( - kv_e8m0.view(batch, t_blocks, kv_block_size, k_tiles, 4) - .permute(0, 1, 3, 4, 2) - .contiguous() - .view(batch * t_blocks, k_tiles, 4, kv_block_size) - # Interleave 4 nts per token group: split [kv_block_size] into - # (NTPW=4, T_per_nt), transpose to (T, NTPW) so 4 consecutive bytes - # per T cover nts 0..3 → 1 dword load. - .view(batch * t_blocks, k_tiles, 4, KVS_NTPW, kv_block_size // KVS_NTPW) - .transpose(-1, -2) - .contiguous() - .view(batch * t_blocks, k_tiles, 4, kv_block_size) + scale_group_size = kv_block_size // KVS_NTPW + kv_e8m0_perm = torch.empty( + batch * t_blocks, + k_tiles, + 4, + kv_block_size, + dtype=torch.uint8, + device=dev, ) + token = torch.arange(kv_block_size, device=dev) + sflat = (token % scale_group_size) * KVS_NTPW + token // scale_group_size + source = kv_e8m0.view(batch * t_blocks, kv_block_size, k_tiles, 4) + kv_e8m0_perm[:, :, :, sflat] = source.permute(0, 2, 3, 1) phys_flat = block_tables.reshape(-1).long() kv_cache = torch.zeros( @@ -766,6 +765,21 @@ def main(): traceback.print_exc() raise + if args.batch == 0 and args.ctx == 0: + test_pa_mqa_logits_fp4_qfp4_kvfp4( + batch=2, + max_ctx=512, + heads=64, + kv_block_size=128, + block_k=256, + num_iters=args.num_iters, + num_warmup=args.num_warmup, + num_warps=args.num_warps, + parallel_unit_num=args.parallel_unit_num, + head_dim=args.head_dim, + bench=False, + ) + if _PERF_SUMMARY: _print_perf_summary() print(" PASS") diff --git a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py index b32d1194d61..9d9ab58f730 100755 --- a/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py +++ b/op_tests/test_flydsl_pa_mqa_logits_fp4_prefill.py @@ -140,7 +140,8 @@ def indexer_k_fp4_paged_preshuffle(k, slot_mapping, kv_cache, kv_scale, kv_block phys = sm // kv_block_size boff = sm % kv_block_size kv_cache[phys, :, :, boff, :] = packed - sflat = (boff % 16) * KVS_NTPW + (boff // 16) + scale_group_size = kv_block_size // KVS_NTPW + sflat = (boff % scale_group_size) * KVS_NTPW + boff // scale_group_size kv_scale[phys, :, :, sflat] = e8m0 return kv_cache, kv_scale @@ -800,6 +801,13 @@ def main(): run_case(3, [[30], [200], [100, 150]], seed=2) run_case(2, [[16, 200], [64, 128]], heads=128, seed=3) run_case(2, [[(10, 50), (64, 200)], [(0, 100), (130, 256)]], seed=4) + run_case( + 2, + [[50, 120, 200], [40, 100]], + kv_block_size=128, + block_k=256, + seed=5, + ) if args.bench: kvb = 64