From f11a6afb365947a5343f00d76851b5ad85378e4a Mon Sep 17 00:00:00 2001 From: who who who <17592563+fsx950223@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:05:42 +0000 Subject: [PATCH] [FlyDSL] Simplify PA decode reducer bound and reuse shared helpers - Fold the reducer's partition_bound into max_partitions: pa_decode passes one tight per-call bound; above a wave the power-of-two round-up is capped by the CU count. - Use dtypes.fp8 for the KV FP8 dtype, tensor_shim._FX_ELEM for the reducer's dtype lookup, and fx.as_ir_value instead of _to_raw. - Name the gfx1250 thread count Gfx1250Traits.BLOCK_THREADS instead of hardcoding 128. - Treat any nonpositive sliding_window as disabled in pa_decode and plan_pa_decode instead of rejecting values below -1. Co-Authored-By: Claude Opus 5.5 --- aiter/ops/flydsl/kernels/pa_decode/op_gemm.py | 13 ++++--- aiter/ops/flydsl/kernels/pa_decode/op_lds.py | 12 +++---- aiter/ops/flydsl/kernels/pa_decode/traits.py | 1 + aiter/ops/flydsl/kernels/pa_decode_kernel.py | 8 ++--- aiter/ops/flydsl/kernels/pa_decode_plan.py | 7 ++-- aiter/ops/flydsl/kernels/pa_decode_reduce.py | 16 +++------ aiter/ops/flydsl/pa_decode.py | 34 ++++++++----------- 7 files changed, 40 insertions(+), 51 deletions(-) diff --git a/aiter/ops/flydsl/kernels/pa_decode/op_gemm.py b/aiter/ops/flydsl/kernels/pa_decode/op_gemm.py index c18cbb7b61b..0f0d7709b91 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/op_gemm.py +++ b/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"), 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"), reuseA=False, reuseB=False, diff --git a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py b/aiter/ops/flydsl/kernels/pa_decode/op_lds.py index b29bb626f3f..a87d3c2992e 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/op_lds.py +++ b/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 diff --git a/aiter/ops/flydsl/kernels/pa_decode/traits.py b/aiter/ops/flydsl/kernels/pa_decode/traits.py index de94a95d83f..0cd5f3d8322 100644 --- a/aiter/ops/flydsl/kernels/pa_decode/traits.py +++ b/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 diff --git a/aiter/ops/flydsl/kernels/pa_decode_kernel.py b/aiter/ops/flydsl/kernels/pa_decode_kernel.py index 74f174fe892..dfb2d8ac4d7 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_kernel.py +++ b/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, ) diff --git a/aiter/ops/flydsl/kernels/pa_decode_plan.py b/aiter/ops/flydsl/kernels/pa_decode_plan.py index 35ef6657fd1..f4f46514af6 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_plan.py +++ b/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") diff --git a/aiter/ops/flydsl/kernels/pa_decode_reduce.py b/aiter/ops/flydsl/kernels/pa_decode_reduce.py index 81fe70a66cd..954c28c9f5c 100644 --- a/aiter/ops/flydsl/kernels/pa_decode_reduce.py +++ b/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 = () diff --git a/aiter/ops/flydsl/pa_decode.py b/aiter/ops/flydsl/pa_decode.py index 4e33a3e3190..700fba7eb8b 100644 --- a/aiter/ops/flydsl/pa_decode.py +++ b/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, )