[FlyDSL] [Misc] Simplify PA decode reducer bound and reuse shared helpers

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

@@ -8,7 +8,6 @@
 from flydsl.expr import const_expr, range_constexpr
 from flydsl.expr.typing import T
 
-from ..tensor_shim import _to_raw
 from .traits import MFMA_ACC_ELEMS
 
 
@@ -106,9 +105,9 @@ def mma(a, b, acc, k, bf16=False):
         if const_expr(bf16):
             result = fx.rocdl.wmma_f32_16x16x32_bf16(
                 fx.Vector.make_type(8, fx.Float32),
-                _to_raw(a),
-                _to_raw(b),
-                _to_raw(acc),
+                fx.as_ir_value(a),
+                fx.as_ir_value(b),
+                fx.as_ir_value(acc),
                 modC=ir.Attribute.parse("#rocdl<wmma_c_modifier none>"),
                 reuseA=False,
                 reuseB=False,
@@ -120,9 +119,9 @@ def mma(a, b, acc, k, bf16=False):
         else:
             result = fx.rocdl.wmma_f32_16x16x64_fp8_fp8(
                 fx.Vector.make_type(8, fx.Float32),
-                _to_raw(a),
-                _to_raw(b),
-                _to_raw(acc),
+                fx.as_ir_value(a),
+                fx.as_ir_value(b),
+                fx.as_ir_value(acc),
                 modC=ir.Attribute.parse("#rocdl<wmma_c_modifier none>"),
                 reuseA=False,
                 reuseB=False,

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

@@ -12,7 +12,7 @@
 
 from .. import dpp_utils
 from ..kernels_common import create_llvm_ptr
-from ..tensor_shim import _to_raw, buf_base_i64, ptr_buf_tensor
+from ..tensor_shim import buf_base_i64, ptr_buf_tensor
 from ..utils import rcp_f32
 from .traits import MFMA_MNK
 
@@ -768,7 +768,7 @@ def fp8(self, values):
         # instruction converts eight FP32 values into two packed dwords.
         return fx.Vector(
             fx.rocdl.cvt_scalef32_pk8_fp8_f32(
-                T.i32x2, _to_raw(values), _to_raw(fx.Float32(1.0))
+                T.i32x2, fx.as_ir_value(values), fx.as_ir_value(fx.Float32(1.0))
             )
         )
 
@@ -1243,8 +1243,8 @@ def stage_vector_dma(self, token_base, kv_off):
         key_base = buf_base_i64(ctx.key_cache_ptr)
         value_base = buf_base_i64(ctx.value_cache_ptr)
         # K cache has contiguous 16-byte head-dimension chunks.
-        for i in range_constexpr(t.TOKENS * t.head_dim // t.chunk // 128):
-            chunk = ctx.tid + i * 128
+        for i in range_constexpr(t.TOKENS * t.head_dim // t.chunk // t.BLOCK_THREADS):
+            chunk = ctx.tid + i * t.BLOCK_THREADS
             token = chunk // (t.head_dim // t.chunk)
             d = (chunk % (t.head_dim // t.chunk)) * t.chunk
             absolute = token_base + token
@@ -1267,8 +1267,8 @@ def stage_vector_dma(self, token_base, kv_off):
             )
 
         # Both supported V layouts have contiguous 16-byte token chunks.
-        for i in range_constexpr(t.head_dim * (t.TOKENS // t.chunk) // 128):
-            chunk = ctx.tid + i * 128
+        for i in range_constexpr(t.head_dim * (t.TOKENS // t.chunk) // t.BLOCK_THREADS):
+            chunk = ctx.tid + i * t.BLOCK_THREADS
             d = chunk // (t.TOKENS // t.chunk)
             token = (chunk % (t.TOKENS // t.chunk)) * t.chunk
             absolute = token_base + token

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

@@ -770,6 +770,7 @@ class Gfx1250Traits:
 
     WAVE = 32
     NWARP = 4
+    BLOCK_THREADS = NWARP * WAVE
     FP8_MAX = 448.0
 
     @property

aiter/ops/flydsl/kernels/pa_decode_kernel.py

@@ -503,7 +503,7 @@ class SharedStorage:
         direct_output=direct_output,
     )
 
-    @flyc.kernel(name=kernel_name, known_block_size=(128, 1, 1))
+    @flyc.kernel(name=kernel_name, known_block_size=(t.BLOCK_THREADS, 1, 1))
     def pa_decode_tile_kernel(
         pmax: fx.Pointer,
         psum: fx.Pointer,
@@ -569,8 +569,8 @@ def pa_decode_tile_kernel(
         elif const_expr(t.direct_output):
             output = fx.recast_iter(t.q_dtype, pout)
             tid = fx.Int32(gpu.thread_id("x"))
-            for i in range_constexpr((t.rows * t.head_dim + 127) // 128):
-                index = i * 128 + tid
+            for i in range_constexpr(-(-t.rows * t.head_dim // t.BLOCK_THREADS)):
+                index = i * t.BLOCK_THREADS + tid
                 if index < t.rows * t.head_dim:
                     row, d = index // t.head_dim, index % t.head_dim
                     kv_query = fx.Int32(gpu.block_id("y"))
@@ -636,7 +636,7 @@ def launch(
                 kv_heads * t.query_splits,
                 1,
             ),
-            block=(128, 1, 1),
+            block=(t.BLOCK_THREADS, 1, 1),
             stream=stream,
         )
 

aiter/ops/flydsl/kernels/pa_decode_plan.py

@@ -216,13 +216,12 @@ def plan_pa_decode(
     must fit the device's per-block shared-memory capacity.
 
     Lengths include MTP tokens. Positive windows include each query's position;
-    0 and -1 disable them. Plans cover the MTP window union: match the window
-    and query length when decoding or refreshing. Dense plans ignore query length.
+    nonpositive windows disable them. Plans cover the MTP window union: match
+    the window and query length when decoding or refreshing. Dense plans ignore
+    query length.
     """
     if not isinstance(sliding_window, int):
         raise TypeError("sliding_window must be an int")
-    if sliding_window < -1:
-        raise ValueError("sliding_window must be -1, 0, or positive")
     sliding_window = max(sliding_window, 0)
     if not isinstance(query_length, int):
         raise TypeError("query_length must be an int")

aiter/ops/flydsl/kernels/pa_decode_reduce.py

@@ -10,15 +10,9 @@
 
 from aiter.jit.utils.chip_info import get_gfx_runtime
 
-from .tensor_shim import buf_base_i64, buf_copy_load, ptr_buf_tensor
+from .tensor_shim import _FX_ELEM, buf_base_i64, buf_copy_load, ptr_buf_tensor
 from .utils import rcp_f32, udiv_const, urem_const
 
-_DTYPE_MAP = {
-    "f32": fx.Float32,
-    "f16": fx.Float16,
-    "bf16": fx.BFloat16,
-}
-
 WAVE = 64
 # Above this partition bound the CTA reduces in LDS; at most DUAL_REGISTER
 # live partitions still take its register branch.
@@ -91,7 +85,7 @@ def compile_pa_decode_ps_reduce(
             f"head_size must be a multiple of 64 in [64, 1024], got {head_size}"
         )
     for dtype_str in (output_dtype_str, logits_dtype_str, sink_dtype_str):
-        if dtype_str not in _DTYPE_MAP:
+        if dtype_str not in ("f32", "f16", "bf16"):
             raise ValueError(f"Unsupported FlyDSL dtype: {dtype_str!r}")
     if query_group_size <= 0:
         raise ValueError(f"query_group_size must be positive, got {query_group_size}")
@@ -110,9 +104,9 @@ def compile_pa_decode_ps_reduce(
     if num_xcds < 0:
         raise ValueError(f"num_xcds must be nonnegative, got {num_xcds}")
 
-    output_dtype = _DTYPE_MAP[output_dtype_str]
-    logits_dtype = _DTYPE_MAP[logits_dtype_str]
-    sink_dtype = _DTYPE_MAP[sink_dtype_str]
+    output_dtype = _FX_ELEM[output_dtype_str]
+    logits_dtype = _FX_ELEM[logits_dtype_str]
+    sink_dtype = _FX_ELEM[sink_dtype_str]
     log2e = 1.4426950408889634
 
     low_rungs = high_rungs = dual_rungs = lds_rungs = ()

aiter/ops/flydsl/pa_decode.py

@@ -15,6 +15,7 @@
 
 from aiter.jit.utils.chip_info import get_gfx_runtime
 from aiter.ops.triton.utils.device_info import get_num_xcds
+from aiter.utility import dtypes
 
 from .kernels.pa_decode_kernel import KV_COMPUTE_BLOCK, compile_pa_decode_tile
 from .kernels.pa_decode_plan import PADecodePlan
@@ -57,12 +58,11 @@ def launch_pa_decode_ps_reduce(
     query_seq_len: int,
     query_group_size: int,
     head_size: int,
-    partition_bound: int,
     max_partitions: int,
     stream: torch.cuda.Stream,
     reduce_info: torch.Tensor,
 ) -> None:
-    """Reduce packed partials; every sequence has at most ``partition_bound``."""
+    """Reduce packed partials; every sequence has at most ``max_partitions``."""
     props = torch.cuda.get_device_properties(output.device)
     if max_partitions > props.multi_processor_count:
         raise ImportError(
@@ -76,10 +76,13 @@ def launch_pa_decode_ps_reduce(
     batch, kv_heads = output.shape[0], output.shape[2]
     rows = query_seq_len * query_group_size
     # Exact bounds keep the register ladder short; above a wave a power of two
-    # bounds the specialization count.
-    partitions = min(partition_bound, max_partitions)
+    # bounds the specialization count. Per-sequence buffer bounds zero-fill
+    # partitions beyond the live count, so rounding past the bound is safe.
+    partitions = max_partitions
     if partitions > register_partitions:
-        partitions = min(max_partitions, 1 << (partitions - 1).bit_length())
+        partitions = min(
+            props.multi_processor_count, 1 << (partitions - 1).bit_length()
+        )
     # A per-sequence span <= INT32_MAX keeps every attempted byte offset,
     # including inactive partitions, below 2**32, so out-of-bounds reads
     # cannot wrap back into valid data.
@@ -221,7 +224,7 @@ def pa_decode(
     MTP lengths include the query tokens and use dense causal masking.
     Independently selected sparse queries need separate table rows and
     query_length=1. Positive ``sliding_window`` includes each query's own
-    position; 0 and -1 disable it.
+    position; nonpositive values disable it.
 
     ``work_plan`` is required and must be built with ``plan_pa_decode`` before
     calling decode. ``pa_decode_tuning.prepare_pa_decode_plan`` uses native
@@ -269,8 +272,6 @@ def pa_decode(
         raise NotImplementedError("pa_decode does not support ALiBi")
     if not isinstance(sliding_window, int):
         raise TypeError("sliding_window must be an int")
-    if sliding_window < -1:
-        raise ValueError("sliding_window must be -1, 0, or positive")
     sliding_window = max(sliding_window, 0)
     if not isinstance(query_length, int):
         raise TypeError("query_length must be an int")
@@ -340,18 +341,13 @@ def pa_decode(
         )
 
     arch = get_gfx_runtime()
-    expected_fp8_dtype = {
-        "gfx942": torch.float8_e4m3fnuz,
-        "gfx950": torch.float8_e4m3fn,
-        "gfx1250": torch.float8_e4m3fn,
-    }.get(arch)
-    if expected_fp8_dtype is None:
+    if arch not in ("gfx942", "gfx950", "gfx1250"):
         raise NotImplementedError(
             f"pa_decode only supports gfx942, gfx950 and gfx1250, got {arch}"
         )
     is_bf16_kv = key_cache.dtype == torch.bfloat16
     kv_dtype = "bf16" if is_bf16_kv else "fp8"
-    expected_kv_dtype = torch.bfloat16 if is_bf16_kv else expected_fp8_dtype
+    expected_kv_dtype = torch.bfloat16 if is_bf16_kv else dtypes.fp8
     compute_type = key_cache.dtype if compute_type is None else compute_type
     if compute_type != expected_kv_dtype:
         raise NotImplementedError(
@@ -743,7 +739,8 @@ def normalize_scale(scale, name):
             num_seqs, query_length, num_kv_heads, query_group_size, head_dim
         )
         # Bound the reducer without shrinking the reusable plan or scratch.
-        partition_bound = min(
+        max_partitions = min(
+            num_partitions,
             work_plan.capacity,
             max(1, -(-schedule_context_length // KV_COMPUTE_BLOCK)),
         )
@@ -752,7 +749,7 @@ def normalize_scale(scale, name):
             window_tiles = (
                 sliding_window + query_length - 2 + KV_COMPUTE_BLOCK - 1
             ) // KV_COMPUTE_BLOCK + 1
-            partition_bound = min(partition_bound, window_tiles)
+            max_partitions = min(max_partitions, window_tiles)
         launch_pa_decode_ps_reduce(
             output_5d,
             psum,
@@ -771,8 +768,7 @@ def normalize_scale(scale, name):
             query_seq_len=query_length,
             query_group_size=query_group_size,
             head_size=head_dim,
-            partition_bound=partition_bound,
-            max_partitions=num_partitions,
+            max_partitions=max_partitions,
             stream=s,
             reduce_info=work_plan.reduce_info,
         )