diff --git a/.github/workflows/mslk_ci_rocm.yml b/.github/workflows/mslk_ci_rocm.yml index 0bab62b8..becfb464 100644 --- a/.github/workflows/mslk_ci_rocm.yml +++ b/.github/workflows/mslk_ci_rocm.yml @@ -37,6 +37,11 @@ on: - 'mslk/attention/flydsl/**' - 'test/attention/flydsl/**' - 'test/flydsl/**' + # FlyDSL paged-attention decode backend (fmha op layer) and tests + - 'mslk/attention/fmha/flydsl/**' + - 'mslk/attention/fmha/flydsl_decoder.py' + - 'mslk/attention/fmha/flydsl_splitk.py' + - 'test/attention/fmha/**' # GEMM tests - 'test/gemm/gemm_test.py' # AMD/ROCm Triton GEMM kernels diff --git a/bench/attn/decoder_bench.py b/bench/attn/decoder_bench.py new file mode 100644 index 00000000..9d5c1e4d --- /dev/null +++ b/bench/attn/decoder_bench.py @@ -0,0 +1,277 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Paged-attention decode benchmark: FlyDSL vs CK vs Triton. + +Benchmarks the fmha decode forward ops (see ``decoder_ops.py``) across decode shapes +and reports latency, achieved HBM bandwidth, and memory-bandwidth utilization. One +row per op per shape; compare rows to read relative speedups. + +``--dtype {bf16,f16}`` runs the dense ops; ``--dtype fp8`` quantizes the KV cache and +runs the fp8 ops (FlyDSLFp8, TritonFp8) instead. + +Usage: + python bench/attn/decoder_bench.py --shapes decode_llm + python bench/attn/decoder_bench.py --shapes sweep_kv --dtype fp8 --export-csv + python bench/attn/decoder_bench.py --kernels FlyDSLDecode,CKDecode,TritonSplitK +""" + +import os +import sys +from dataclasses import dataclass +from datetime import datetime +from typing import Any, Callable, Optional + +import click +import torch +import triton # @manual=//triton:triton +from mslk.bench.attn.decoder_ops import DecodeOpBase, get_decode_ops +from mslk.bench.common.utils import BenchOptions, common_bench_options, profiler + +ShapeList = list[tuple[int, int, int, int, int]] + +shape_registry: dict[str, Callable[[], ShapeList]] = {} + + +def register_shapes( + name: str, +) -> Callable[[Callable[[], ShapeList]], Callable[[], ShapeList]]: + def decorator(fn: Callable[[], ShapeList]) -> Callable[[], ShapeList]: + shape_registry[name] = fn + return fn + + return decorator + + +@register_shapes("decode_llm") +def _shapes_decode_llm() -> ShapeList: + """Common LLM decode shapes with various GQA ratios: (B, Hq, Hkv, kv_len, D).""" + return [ + (1, 32, 8, 512, 128), + (1, 32, 8, 2048, 128), + (1, 32, 8, 4096, 128), + (8, 32, 8, 2048, 128), + (16, 32, 8, 2048, 128), + (1, 64, 8, 2048, 128), + (1, 64, 16, 2048, 128), + (1, 128, 16, 2048, 128), + (1, 32, 4, 2048, 256), + (8, 32, 4, 2048, 256), + ] + + +@register_shapes("sweep_kv") +def _shapes_sweep_kv() -> ShapeList: + """Sweep KV sequence length.""" + return [ + (1, 32, 8, kv_len, 128) + for kv_len in [128, 256, 512, 1024, 2048, 4096, 8192, 16384] + ] + + +@register_shapes("sweep_batch") +def _shapes_sweep_batch() -> ShapeList: + """Sweep batch size.""" + return [(B, 32, 8, 2048, 128) for B in [1, 2, 4, 8, 16, 32, 64]] + + +def _bytes_read_write( + B: int, Hq: int, Hkv: int, kv_seqlen: int, D: int, dtype: str +) -> int: + """Approximate HBM traffic for one decode step (bytes): Q + K + V + output. + + Query/output are 16-bit; the KV cache is 1 byte under fp8, else 16-bit. + """ + io_elem = 2 # bf16/f16 query + output + kv_elem = 1 if dtype == "fp8" else 2 + q_read = B * Hq * D * io_elem + kv_read = B * kv_seqlen * Hkv * D * kv_elem * 2 # K and V + out_write = B * Hq * D * io_elem + return q_read + kv_read + out_write + + +@dataclass +class Metrics: + op: str + B: int + Hq: int + Hkv: int + kv_seqlen: int + D: int + dtype: str + ms: float = 0.0 + gbps: float = 0.0 + mem_bw_util: float = 0.0 + + @staticmethod + def header() -> str: + header = ( + f"{'OpName':<16} {'B':>4} {'Hq':>4} {'Hkv':>4} {'KV':>6} {'D':>4} " + f"{'dtype':>6} | {'Ms':>10} {'GB/s':>10} {'Mem BW Util %':>14}" + ) + divider = "-" * len(header) + return f"Decoder Attention Bench\n{divider}\n{header}\n{divider}" + + def __str__(self) -> str: + return ( + f"{self.op:<16} {self.B:>4} {self.Hq:>4} {self.Hkv:>4} {self.kv_seqlen:>6} " + f"{self.D:>4} {self.dtype:>6} | {self.ms:>10.3f} {self.gbps:>10.2f} " + f"{self.mem_bw_util:>14.2f}" + ) + + def as_dict(self) -> dict[str, Any]: + return { + "op": self.op, + "B": self.B, + "Hq": self.Hq, + "Hkv": self.Hkv, + "kv_seqlen": self.kv_seqlen, + "D": self.D, + "dtype": self.dtype, + "ms": self.ms, + "gbps": self.gbps, + "mem_bw_util": self.mem_bw_util, + } + + +def benchmark( + ops: list[DecodeOpBase], + B: int, + Hq: int, + Hkv: int, + kv_seqlen: int, + D: int, + dtype: torch.dtype, + dtype_str: str, + mem_bw_roofline_gbps: float, + opts: BenchOptions, +) -> list[Metrics]: + """Benchmark every op for one decode shape.""" + nbytes = _bytes_read_write(B, Hq, Hkv, kv_seqlen, D, dtype_str) + results: list[Metrics] = [] + for op in ops: + shape_str = f"(B={B}, Hq={Hq}, Hkv={Hkv}, KV={kv_seqlen}, D={D})" + print(f"Benchmarking {op.name} with {shape_str}") + try: + args = op.setup(B, Hq, Hkv, kv_seqlen, D, dtype) + op.compute(*args) # warmup / sanity + except Exception as e: + print(f"Decode op {op.name} failed to run due to error: {e}.") + continue + with profiler(enabled=opts.trace, with_stack=True): + ms = op.benchmark(*args, opts=opts) + gbps = nbytes / (ms / 1e3) / 1e9 + results.append( + Metrics( + op=op.name, + B=B, + Hq=Hq, + Hkv=Hkv, + kv_seqlen=kv_seqlen, + D=D, + dtype=dtype_str, + ms=ms, + gbps=gbps, + mem_bw_util=(gbps / mem_bw_roofline_gbps) * 100, + ) + ) + return results + + +def collect_ops(kernels: Optional[list[str]], dtype: str) -> list[DecodeOpBase]: + ops = [ + op for op in get_decode_ops() if op.supported and dtype in op.supported_dtypes + ] + if kernels is None: + return ops + return [op for op in ops if op.name in kernels] + + +@click.command() +@common_bench_options(shape_registry) +@click.option( + "--dtype", + default="bf16", + type=click.Choice(["bf16", "f16", "fp8"]), + help="KV-cache dtype. fp8 quantizes the KV cache (bf16 query) and runs the fp8 ops.", +) +def invoke_main( + output_dir: str, + export_csv: bool, + kernels: Optional[str], + cuda_graph: bool, + rotating_buffer: bool, + shapes: Optional[str], + trace: bool, + rep_ms: int, + dtype: str, +) -> None: + # fp8 uses a bf16 query with a quantized KV cache (see the fp8 ops). + torch_dtype = { + "bf16": torch.bfloat16, + "f16": torch.float16, + "fp8": torch.bfloat16, + }[dtype] + + kernel_filter = kernels.strip().split(",") if kernels else None + ops = collect_ops(kernel_filter, dtype) + if not ops: + available = ", ".join(op.name for op in get_decode_ops()) + print(f"No matching supported ops. Available: {available}.") + sys.exit(1) + + if shapes: + if shapes not in shape_registry: + print( + f"Shape '{shapes}' not found. Valid: {', '.join(shape_registry.keys())}." + ) + sys.exit(1) + shape_list = shape_registry[shapes]() + else: + shape_list = shape_registry["decode_llm"]() + + opts = BenchOptions( + cuda_graph=cuda_graph, + rotating_buffer=rotating_buffer, + rep_ms=rep_ms, + trace=trace, + ) + + mem_bw_gbps = triton.testing.get_dram_gbps() + results: list[Metrics] = [] + for B, Hq, Hkv, kv_seqlen, D in shape_list: + results.extend( + benchmark( + ops, B, Hq, Hkv, kv_seqlen, D, torch_dtype, dtype, mem_bw_gbps, opts + ) + ) + + print("") + print(Metrics.header()) + for m in results: + print(m) + + print("") + print(f"Hardware: {torch.cuda.get_device_name()}") + print(f" Memory BW: {mem_bw_gbps:.2f} GB/s") + print("") + print("Benchmark Settings:") + print(f" CUDA graph: {cuda_graph}") + print(f" Buffer rotation: {rotating_buffer}") + print(f" dtype: {dtype}") + + if export_csv: + import pandas as pd + + os.makedirs(output_dir, exist_ok=True) + datetime_str = datetime.now().strftime("%Y%m%d_%H%M%S") + csv_file = os.path.join(output_dir, f"decoder_bench_{datetime_str}.csv") + pd.DataFrame([m.as_dict() for m in results]).to_csv(csv_file, index=False) + print(f"CSV saved to {csv_file}") + + +if __name__ == "__main__": + invoke_main() # pragma: no cover diff --git a/bench/attn/decoder_ops.py b/bench/attn/decoder_ops.py new file mode 100644 index 00000000..9a91d8c9 --- /dev/null +++ b/bench/attn/decoder_ops.py @@ -0,0 +1,341 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Paged-attention decode ops for decoder_bench. + +Each op wraps an fmha forward operator (FlyDSL / CK / Triton) and is benchmarked +through the standard ``op.apply`` interface, so all backends share one code path +and input layout. Mirrors the ops-registry pattern of ``bench/gemm/gemm_ops.py``. +""" + +import abc + +import torch +from mslk.attention.fmha import ( + ck_decoder, + ck_splitk, + flydsl_decoder, + flydsl_splitk, + triton_splitk, +) +from mslk.attention.fmha.attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask +from mslk.attention.fmha.common import Inputs, InputsFp8 +from mslk.bench.common.utils import BenchOptions, do_bench +from mslk.utils.triton.fp8_utils import get_fp8_constants + +try: + from mslk.attention.fmha.flydsl.fp8_paged_cache import dense_kv_to_fp8_paged + from mslk.attention.fmha.flydsl.pa_decode_fp8 import ( + get_recommended_splits, + KV_COMPUTE_BLOCK, + pa_decode_ps_launch, + ) + from mslk.attention.fmha.flydsl.pa_decode_fp8_dispatch import ( + is_fp8_paged_decode_available, + ) + from mslk.flydsl.common import is_flydsl_available + + FLYDSL_ENABLED = True +except ImportError: + FLYDSL_ENABLED = False + +DENSE_DTYPES = ("bf16", "f16") + + +decode_op_registry: list["DecodeOpBase"] = [] + + +def register_decode_op(op): + """Decorator that registers a single instance of a decode op.""" + decode_op_registry.append(op()) + return op + + +def get_decode_ops() -> list["DecodeOpBase"]: + """Return all registered decode ops.""" + return decode_op_registry + + +class DecodeOpBase(metaclass=abc.ABCMeta): + """A paged-decode operator benchmarked through the fmha forward interface.""" + + # fmha forward op class (subclass of AttentionFwOpBase). + OP = None + + @property + def name(self) -> str: + return self.__class__.__name__ + + # dtypes (``--dtype``) this op benchmarks. Dense ops handle bf16/f16; fp8 ops + # quantize the KV cache and only run under ``--dtype fp8``. + supported_dtypes: tuple[str, ...] = DENSE_DTYPES + + @property + def supported(self) -> bool: + """Whether this op can run on the current device/build.""" + return self.OP is not None and self.OP.is_available() + + def setup( + self, + B: int, + Hq: int, + Hkv: int, + kv_seqlen: int, + D: int, + dtype: torch.dtype, + ) -> tuple: + """Build the decode inputs (Q, K, V, attn_bias) for one shape. + + Canonical BMGHK decode layout (matches ``_test_decoder`` in the test suite): + ``G = Hkv`` groups, ``Hq // Hkv`` query heads per group. K/V are stored folded + (one head per group) and broadcast to the query-head count with ``.expand()``, + i.e. stride-0 — the realistic GQA cache. Q/K/V are returned as plain tensors so + the shared ``do_bench`` rotating buffer can size and rotate them. + """ + dev = "cuda" + Hpg = Hq // Hkv + shape = (1, B * kv_seqlen, Hkv, Hpg, D) + q = torch.randn(1, B, Hkv, Hpg, D, dtype=dtype, device=dev) + k = torch.randn(1, B * kv_seqlen, Hkv, 1, D, dtype=dtype, device=dev).expand( + shape + ) + v = torch.randn(1, B * kv_seqlen, Hkv, 1, D, dtype=dtype, device=dev).expand( + shape + ) + attn_bias = BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens( + q_seqlen=[1] * B, + kv_seqlen=[kv_seqlen] * B, + kv_padding=kv_seqlen, + ) + attn_bias.k_seqinfo.to(dev) + attn_bias.q_seqinfo.to(dev) + return q, k, v, attn_bias + + def compute(self, q, k, v, attn_bias) -> torch.Tensor: + inp = Inputs(q, k, v, attn_bias=attn_bias, scale=float(q.shape[-1] ** -0.5)) + out, _ = self.OP.apply(inp, False) + return out + + def benchmark(self, *args, opts: BenchOptions) -> float: + """Benchmark runtime (ms) of this op.""" + return do_bench(lambda *a: self.compute(*a), args, opts) + + +class _FlyDSLDecodeOp(DecodeOpBase): + """FlyDSL ops additionally require the flydsl package + a supported arch.""" + + @property + def supported(self) -> bool: + return ( + FLYDSL_ENABLED + and is_flydsl_available() + and self.OP is not None + and self.OP.is_available() + ) + + +@register_decode_op +class FlyDSLDecode(_FlyDSLDecodeOp): + OP = flydsl_decoder.FwOp + + +@register_decode_op +class FlyDSLSplitK(_FlyDSLDecodeOp): + OP = flydsl_splitk.FwOp + + +@register_decode_op +class CKDecode(DecodeOpBase): + OP = ck_decoder.FwOp + + +@register_decode_op +class CKSplitK(DecodeOpBase): + OP = ck_splitk.FwOp + + +@register_decode_op +class TritonSplitK(DecodeOpBase): + OP = triton_splitk.FwOp + + def setup(self, B, Hq, Hkv, kv_seqlen, D, dtype) -> tuple: + # Triton split-K does not broadcast KV heads itself; materialize the expanded + # (stride-0) K/V into contiguous tensors. + q, k, v, attn_bias = super().setup(B, Hq, Hkv, kv_seqlen, D, dtype) + return q, k.contiguous(), v.contiguous(), attn_bias + + +# --------------------------------------------------------------------------- # +# fp8 KV-cache ops (--dtype fp8): the KV cache is quantized once in setup, and +# only these ops run under fp8, so the dense and fp8 FlyDSL kernels never share a +# process. +# --------------------------------------------------------------------------- # + + +def _quant_pack_triton_fp8(x: torch.Tensor): + """Quantize dense KV to Triton's int32-packed asymmetric fp8 format. + Returns (packed_int32, scale_shift_int32) as ``InputsFp8`` expects.""" + fp8_dtype = get_fp8_constants()[0] + fmax = torch.finfo(fp8_dtype).max + Bx, M, G, H, Dx = x.shape + xr = x.reshape(-1, Dx).float() + shift = xr.mean(-1) + xc = xr - shift[..., None] + s = torch.nan_to_num(xc.abs().max(-1)[0] / fmax, posinf=1) + xq = (xc / s[..., None]).to(fp8_dtype) + packed = xq.view(torch.uint8).reshape(Bx, M, G, H, Dx).view(torch.int32) + ss = ( + torch.concat( + [s.reshape(Bx, M, G, H, 1).half(), shift.reshape(Bx, M, G, H, 1).half()], + dim=-1, + ) + .flatten(-2) + .view(torch.int32) + ) + return packed, ss + + +@register_decode_op +class FlyDSLFp8(DecodeOpBase): + """FlyDSL native-fp8 (e4m3fn, per-token) paged decode over a precomputed fp8 cache. + + Runs the kernel directly (there is no dense fmha op for native fp8). Scratch is + preallocated in setup so the launch is CUDA-graph capturable. + """ + + supported_dtypes = ("fp8",) + + @property + def supported(self) -> bool: + return FLYDSL_ENABLED and is_fp8_paged_decode_available() + + def setup(self, B, Hq, Hkv, kv_seqlen, D, dtype) -> tuple: + dev = "cuda" + # Native fp8 path uses G=1 with the KV heads in the H slot (the layout + # dense_kv_to_fp8_paged expects); query dtype is bf16/f16. + q = torch.randn(B, 1, 1, Hq, D, dtype=dtype, device=dev) + k = torch.randn(B, kv_seqlen, 1, Hkv, D, dtype=dtype, device=dev) + v = torch.randn(B, kv_seqlen, 1, Hkv, D, dtype=dtype, device=dev) + key_cache, value_cache, key_scale, value_scale, block_tables = ( + dense_kv_to_fp8_paged(k, v, block_size=16) + ) + out = torch.zeros(B, Hq, D, dtype=q.dtype, device=dev) + q_flat = q.reshape(B, Hq, D).contiguous() + context_lengths = torch.full((B,), kv_seqlen, dtype=torch.int32, device=dev) + num_kv_heads = key_cache.shape[1] + eqgs = Hq // num_kv_heads # query_length == 1 for decode + block_size = key_cache.shape[-2] + mcpn = get_recommended_splits( + B, num_kv_heads, split_kv_blocks=KV_COMPUTE_BLOCK // block_size + ) + exp_sums = torch.zeros( + B, num_kv_heads, mcpn, eqgs, device=dev, dtype=torch.float32 + ) + max_logits = torch.full( + (B, num_kv_heads, mcpn, eqgs), + float("-inf"), + device=dev, + dtype=torch.float32, + ) + tmp_out = torch.zeros( + B, num_kv_heads, mcpn, eqgs, D, device=dev, dtype=torch.bfloat16 + ) + scale = float(D**-0.5) + return ( + out, + q_flat, + key_cache, + value_cache, + context_lengths, + key_scale, + value_scale, + block_tables, + exp_sums, + max_logits, + tmp_out, + mcpn, + scale, + ) + + def compute( + self, + out, + q_flat, + key_cache, + value_cache, + context_lengths, + key_scale, + value_scale, + block_tables, + exp_sums, + max_logits, + tmp_out, + mcpn, + scale, + ) -> torch.Tensor: + pa_decode_ps_launch( + out, + q_flat, + key_cache, + value_cache, + context_lengths, + scale, + key_scale=key_scale, + value_scale=value_scale, + block_tables=block_tables, + max_context_partition_num=mcpn, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=tmp_out, + ) + return out + + +@register_decode_op +class TritonFp8(DecodeOpBase): + """Triton split-K decode over int32-packed asymmetric fp8 KV (``InputsFp8``).""" + + OP = triton_splitk.FwOp + supported_dtypes = ("fp8",) + + def setup(self, B, Hq, Hkv, kv_seqlen, D, dtype) -> tuple: + dev = "cuda" + Hpg = Hq // Hkv + shape = (1, B * kv_seqlen, Hkv, Hpg, D) + q = torch.randn(1, B, Hkv, Hpg, D, dtype=dtype, device=dev) + k = ( + torch.randn(1, B * kv_seqlen, Hkv, 1, D, dtype=dtype, device=dev) + .expand(shape) + .contiguous() + ) + v = ( + torch.randn(1, B * kv_seqlen, Hkv, 1, D, dtype=dtype, device=dev) + .expand(shape) + .contiguous() + ) + attn_bias = BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens( + q_seqlen=[1] * B, + kv_seqlen=[kv_seqlen] * B, + kv_padding=kv_seqlen, + ) + attn_bias.k_seqinfo.to(dev) + attn_bias.q_seqinfo.to(dev) + ki, ks = _quant_pack_triton_fp8(k) + vi, vs = _quant_pack_triton_fp8(v) + return q, ki, vi, ks, vs, attn_bias, float(D**-0.5) + + def compute(self, q, ki, vi, ks, vs, attn_bias, scale) -> torch.Tensor: + inp = InputsFp8( + q, + ki, + vi, + attn_bias=attn_bias, + scale=scale, + k_fp8_scale_shift=ks, + v_fp8_scale_shift=vs, + ) + out, _ = self.OP.apply(inp, False) + return out diff --git a/mslk/attention/fmha/__init__.py b/mslk/attention/fmha/__init__.py index e1344961..b67ac609 100644 --- a/mslk/attention/fmha/__init__.py +++ b/mslk/attention/fmha/__init__.py @@ -21,6 +21,8 @@ flash, flash3, flash_mtia, + flydsl_decoder, + flydsl_splitk, triton_splitk, ) from .attn_bias import ( @@ -58,6 +60,8 @@ MemoryEfficientAttentionCkOp = (ck.FwOp, ck.BwOp) MemoryEfficientAttentionCkDecoderOp = (ck_decoder.FwOp, ck.BwOp) MemoryEfficientAttentionSplitKCkOp = (ck_splitk.FwOp, ck.BwOp) +MemoryEfficientAttentionFlyDSLDecoderOp = (flydsl_decoder.FwOp, ck.BwOp) +MemoryEfficientAttentionSplitKFlyDSLOp = (flydsl_splitk.FwOp, ck.BwOp) MemoryEfficientAttentionCuteFlashAttentionOp = ( cute_blackwell.FwOp, cute_blackwell.BwOp, @@ -981,6 +985,8 @@ def merge_attentions( # noqa: C901 "MemoryEfficientAttentionFlashMtiaAttentionOp", "memory_efficient_attention", "MemoryEfficientAttentionCkOp", + "MemoryEfficientAttentionFlyDSLDecoderOp", + "MemoryEfficientAttentionSplitKFlyDSLOp", "MemoryEfficientAttentionCkDecoderOp", "ALL_FW_OPS", "ALL_BW_OPS", diff --git a/mslk/attention/fmha/_triton/splitk_kernels.py b/mslk/attention/fmha/_triton/splitk_kernels.py index 134ef934..73457e53 100644 --- a/mslk/attention/fmha/_triton/splitk_kernels.py +++ b/mslk/attention/fmha/_triton/splitk_kernels.py @@ -121,6 +121,7 @@ def _fwd_kernel_splitK( # noqa: C901 HAS_ADDITIVE_BIAS: tl.constexpr, NUM_PROGRAMS_DIM2_CONST: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -539,6 +540,7 @@ def _fwd_kernel_splitK( # noqa: C901 # pyrefly: ignore [bad-argument-type] i, IS_HIP, + FP8_FNUZ, QUANTIZE_PV_TO_FP8, QUANTIZE_QK_TO_FP8, USE_FP32_SCALES, @@ -780,10 +782,18 @@ def autotune_kernel(kernel: Callable): if block_n >= block_m ] + # HIP graph capture faulted (HSA_INVALID_PACKET / GPU faults) during autotuning of + # this kernel on gfx950 with ROCm < 7.14; fixed in ROCm 7.14. Disable graph-based + # autotuning only on that affected stack; gfx942, CUDA, and gfx950 on ROCm >= 7.14 + # keep it. + from mslk.utils.device import is_gfx950, rocm_version_at_least + + graph_autotune_broken = is_gfx950() and not rocm_version_at_least(7, 14) + kernel = triton.autotune( configs=TRITON_CONFIGS, key=AUTOTUNER_KEY, - use_cuda_graph=True, + use_cuda_graph=not graph_autotune_broken, prune_configs_by={ "early_config_prune": early_config_prune, }, @@ -834,6 +844,7 @@ def load_dequantize_k_v_group( v_dtype: tl.constexpr, # Q.dtype.element_ty group_id: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -874,6 +885,7 @@ def load_dequantize_k_v_group( q_dtype, v_dtype, IS_HIP, + FP8_FNUZ, QUANTIZE_PV_TO_FP8, QUANTIZE_QK_TO_FP8, USE_FP32_SCALES, @@ -921,6 +933,7 @@ def _process_fp8_quantization( q_dtype: tl.constexpr, v_dtype: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -938,6 +951,7 @@ def _process_fp8_quantization( v_shift if not USE_FP32_SCALES else None, PACKED_PER_VAL, IS_HIP, + FP8_FNUZ, USE_FP32_SCALES, ).to(v_dtype) else: @@ -951,7 +965,9 @@ def _process_fp8_quantization( k_scale, k_shift = _extract_scale_shift(k_scale_shift, IS_HIP, USE_FP32_SCALES) if IS_HIP: if not QUANTIZE_QK_TO_FP8: - k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(q_dtype) + k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, FP8_FNUZ).to( + q_dtype + ) else: # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8 tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") @@ -964,6 +980,7 @@ def _process_fp8_quantization( tl.trans(k_shift) if not USE_FP32_SCALES else None, PACKED_PER_VAL, IS_HIP, + FP8_FNUZ, USE_FP32_SCALES, ).to(q_dtype) k = tl.trans(k_t) @@ -971,7 +988,7 @@ def _process_fp8_quantization( # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8 tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") k_t = tl.trans(k) - k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP) + k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP, FP8_FNUZ) k = tl.trans(k_t) return k, v, v_scale, k_scale @@ -991,7 +1008,9 @@ def _extract_scale_shift( @triton.jit -def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr): +def _unpack_fp8_tensor( + x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr, FP8_FNUZ: tl.constexpr +): """Unpack FP8 K/V tensor from int32 packed representation.""" tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") @@ -1006,8 +1025,9 @@ def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr): unpacked_values, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL) ) - # Convert to FP8 through bitcast - fp8_type = tl.float8e4b8 if IS_HIP else tl.float8e4nv + # Convert to FP8 through bitcast. gfx942 uses e4m3fnuz (float8e4b8); gfx950 and + # CUDA use OCP e4m3fn (float8e4nv). FP8_FNUZ carries the arch decision. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv x_ = unpacked_values.to(tl.uint8).to(fp8_type, bitcast=True) return x_ @@ -1039,18 +1059,20 @@ def _process_int4_quantization( if IS_HIP: k_scale, k_shift = cast_uint32_to_float(k_scale_shift) v_scale, v_shift = cast_uint32_to_float(v_scale_shift) - v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype) - k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype) + # int4 path never reaches the fp8 branch inside dequantize; FP8_FNUZ is unused. + v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype) + k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, False).to(dtype) else: k_scale, k_shift = cast_uint32_to_half2(k_scale_shift) v_scale, v_shift = cast_uint32_to_half2(v_scale_shift) - v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype) + v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype) k_t = dequantize( tl.trans(k), tl.trans(k_scale), tl.trans(k_shift), PACKED_PER_VAL, IS_HIP, + False, ).to(dtype) k = tl.trans(k_t) @@ -1091,6 +1113,7 @@ def dequantize_k_hip( scale, shift, PACKED_PER_VAL: tl.constexpr, + FP8_FNUZ: tl.constexpr, ): """PACKED_PER_VAL is the number of values packed into each element x_. For example, for int4 quantization and x_ of type int32, PACKED_PER_VAL is 8. @@ -1110,8 +1133,8 @@ def dequantize_k_hip( ) if PACKED_PER_VAL == 4: - # FP8 quantization. - fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv + # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv dequant = ( quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale + shift @@ -1140,6 +1163,7 @@ def dequantize( shift, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, # pyrefly: ignore [bad-function-definition] USE_FP32_SCALES: tl.constexpr = False, ): @@ -1160,8 +1184,8 @@ def dequantize( quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL) ) if PACKED_PER_VAL == 4: - # FP8 quantization. - fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv + # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv dequant = ( quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale ) diff --git a/mslk/attention/fmha/flydsl/__init__.py b/mslk/attention/fmha/flydsl/__init__.py new file mode 100644 index 00000000..2e41cd71 --- /dev/null +++ b/mslk/attention/fmha/flydsl/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/mslk/attention/fmha/flydsl/fp8_paged_cache.py b/mslk/attention/fmha/flydsl/fp8_paged_cache.py new file mode 100644 index 00000000..3df191a9 --- /dev/null +++ b/mslk/attention/fmha/flydsl/fp8_paged_cache.py @@ -0,0 +1,113 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Build a native-fp8 paged KV cache (with precomputed scales) from dense KV. + +Quantizes + pages a dense padded f16/bf16 KV cache into the fp8 decode kernel's +layout. Intended to construct a persistent fp8 cache once (e.g. for tests and +benchmarks); it is NOT called per decode step. head_dim % 16 == 0, gfx950. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +_FP8_DTYPE = torch.float8_e4m3fn # OCP e4m3fn — the correct gfx950 fp8 format. +_ELEMS_PER_VEC = 16 # 16 fp8 bytes per 128-bit vector (kernel's head-dim packing). + + +def _pertoken_quant_symmetric( + x: torch.Tensor, fp8_dtype: torch.dtype = _FP8_DTYPE +) -> Tuple[torch.Tensor, torch.Tensor]: + """Symmetric per-token (last-dim) fp8 quant. Returns (xq_fp8, scale_f32).""" + fmax = torch.finfo(fp8_dtype).max + amax = x.abs().amax(dim=-1, keepdim=True).clamp_min(1e-12).to(torch.float32) + scale = amax / fmax + xq = (x.to(torch.float32) / scale).clamp(-fmax, fmax).to(fp8_dtype) + return xq, scale + + +def dense_kv_to_fp8_paged( + key: torch.Tensor, # [B, padding, G, Hkv, D] f16/bf16 + value: torch.Tensor, # [B, padding, G, Hkv, D] f16/bf16 + block_size: int = 16, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize + page a dense per-batch KV cache into the fp8 kernel's layout. + + Blocks are packed batch-major and contiguous, so `block_tables` is the identity. + + Returns (key_cache, value_cache_shuffled, key_scale, value_scale, block_tables): + * key_cache : [num_blocks, Hkv, D//16, block_size, 16] fp8 + * value_cache_shuffled : [num_blocks, Hkv, block_size//16, D, 16] fp8 (trans_v) + * key_scale/value_scale: [num_blocks, Hkv, block_size, 1] f32 (per-token) + * block_tables : [B, blocks_per_seq] int32 (identity, contiguous) + """ + B, padding, G, Hkv, D = key.shape + assert D % _ELEMS_PER_VEC == 0, f"head_dim {D} must be a multiple of 16" + assert padding % block_size == 0, ( + f"padding {padding} must be a multiple of block_size {block_size}" + ) + # GQA: fold (B, G) into the batch/seq axis (one sequence per KV head); query folds + # the same way so group g's query heads pair with KV group g. + from .pa_decode_fp8 import KV_COMPUTE_BLOCK + + BG = B * G + dev = key.device + + # GOTCHA: kernel reads KV_COMPUTE_BLOCK // block_size block-table entries per + # partition; a shorter context would read out of bounds -> GPU fault. Pad each + # seq's block count up to one full partition (extras masked by context_lengths). + min_blocks_per_seq = KV_COMPUTE_BLOCK // block_size + blocks_per_seq = max(padding // block_size, min_blocks_per_seq) + padded = blocks_per_seq * block_size + num_blocks = BG * blocks_per_seq + + # [B, padding, G, Hkv, D] -> [B*G, padding, Hkv, D]. + kbg = key.permute(0, 2, 1, 3, 4).reshape(BG, padding, Hkv, D) + vbg = value.permute(0, 2, 1, 3, 4).reshape(BG, padding, Hkv, D) + if padded != padding: + # GOTCHA: pad with ONES not zeros. Zeros quantize to a ~0 scale and can + # dequant to inf/NaN before context_lengths masks them out. + pad_k = kbg.new_ones(BG, padded - padding, Hkv, D) + kbg = torch.cat([kbg, pad_k], dim=1) + vbg = torch.cat([vbg, pad_k], dim=1) + k = kbg.reshape(num_blocks, block_size, Hkv, D).permute(0, 2, 1, 3).contiguous() + v = vbg.reshape(num_blocks, block_size, Hkv, D).permute(0, 2, 1, 3).contiguous() + + # Per-token symmetric quant over D. + qk, ks = _pertoken_quant_symmetric(k) # qk [nb,Hkv,bs,D], ks [nb,Hkv,bs,1] + qv, vs = _pertoken_quant_symmetric(v) + + # Key cache layout: [num_blocks, Hkv, D//16, block_size, 16]. + key_cache = ( + qk.view(num_blocks, Hkv, block_size, D // _ELEMS_PER_VEC, _ELEMS_PER_VEC) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + # Value cache: first [num_blocks, Hkv, D, block_size], then trans_v shuffle to 5D. + qv_t = qv.permute(0, 1, 3, 2).contiguous() # [nb, Hkv, D, bs] + value_cache = ( + qv_t.view(num_blocks, Hkv, D, block_size // _ELEMS_PER_VEC, _ELEMS_PER_VEC) + .permute(0, 1, 3, 2, 4) + .contiguous() + ) + + # Scales must be the kernel's [num_blocks, Hkv, block_size, 1] layout, contiguous + # (strides (Hkv*bs, bs, 1, 1)). + key_scale = ks.contiguous() + value_scale = vs.contiguous() + + # Identity block_tables: folded seq bg (= b*G + g) owns blocks [bg*bps, (bg+1)*bps). + block_tables = ( + torch.arange(num_blocks, dtype=torch.int32, device=dev) + .view(BG, blocks_per_seq) + .contiguous() + ) + return key_cache, value_cache, key_scale, value_scale, block_tables diff --git a/mslk/attention/fmha/flydsl/layout_utils.py b/mslk/attention/fmha/flydsl/layout_utils.py new file mode 100644 index 00000000..095b437a --- /dev/null +++ b/mslk/attention/fmha/flydsl/layout_utils.py @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Layout utilities: dense KV <-> FlyDSL kernel API adaptation. + +Kernel expects 5D BMGHK: Q [B, 1, G, H_q, D], K/V [B, KVMAX, G, H_kv, D] +(H_kv may = 1 for MQA), seq [B] int32. +""" + +from typing import Optional, Tuple + +import torch + + +def canonicalize_qkv_5d( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (Q, K, V) in [B, *, G, H, D] 5D form (4D BMHK promoted with G=1).""" + if Q.ndim == 4: + Q = Q.unsqueeze(2) + if K.ndim == 4: + K = K.unsqueeze(2) + if V.ndim == 4: + V = V.unsqueeze(2) + + assert Q.ndim == 5 and K.ndim == 5 and V.ndim == 5, ( + f"Expected 5D tensors after promotion; got Q={Q.shape}, K={K.shape}" + ) + + # Multiquery (stride-0 H) is handled by the kernel from .stride() directly. + return Q, K, V + + +def normalize_seq_positions( + seq_kv_lens: Optional[torch.Tensor], + B: int, + KV_MAX: int, + device: torch.device, +) -> torch.Tensor: + """Return a [B] int32 tensor of valid KV lengths. + + If ``seq_kv_lens`` is None, all entries are set to ``KV_MAX``. + """ + if seq_kv_lens is None: + return torch.full((B,), KV_MAX, dtype=torch.int32, device=device) + if seq_kv_lens.dtype != torch.int32: + seq_kv_lens = seq_kv_lens.to(torch.int32) + if seq_kv_lens.device != device: + seq_kv_lens = seq_kv_lens.to(device) + return seq_kv_lens.contiguous() + + +def get_split_k_heuristic(B: int, H: int, Mk: int) -> int: + """Mirror of flydsl_splitk.FwOp.get_split_k — used as default split count.""" + bh = max(B * H, 1) + split_k = max(Mk, 1024) // bh + max_chunk_size = 64 if Mk <= 512 and bh <= 64 else 128 + while split_k > 0 and Mk / split_k < max_chunk_size: + split_k = split_k // 2 + split_k = min(split_k, 64) + split_k = max(split_k, 1) + return split_k diff --git a/mslk/attention/fmha/flydsl/pa_decode_dense.py b/mslk/attention/fmha/flydsl/pa_decode_dense.py new file mode 100644 index 00000000..d657b5fb --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_dense.py @@ -0,0 +1,178 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL decode dispatcher — public entry point for the decoder ops. + +Targets gfx942 (CDNA3/MI300) and gfx950 (CDNA4/MI355), wave64. Compute lives in +pa_decode_gfx950 (head-packed fast path, GQA ratio 1..16), pa_decode_gfx950_coop +(per-head fallback), and pa_decode_generic (arch-generic fallback, off-gfx950). +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + +import torch + +from .utils import WARP_SIZE + +NUM_WARPS = 4 +BLOCK_SIZE = NUM_WARPS * WARP_SIZE # 256 + +_CU_COUNT: Optional[int] = None + + +def _get_cu_count() -> int: + global _CU_COUNT + if _CU_COUNT is None: + try: + prop = torch.cuda.get_device_properties(0) + _CU_COUNT = prop.multi_processor_count + except Exception: + _CU_COUNT = 120 # conservative default + return _CU_COUNT + + +def auto_split_k( + B: int, G: int, H_q: int, KV_MAX: int, num_warps: int = NUM_WARPS +) -> int: + """Default split_k for the generic fallback: target ~4 waves to hide memory latency.""" + n_cus = _get_cu_count() + target_ctas = n_cus * 4 # 4 waves + base_ctas = B * G * H_q + if base_ctas >= target_ctas: + return 1 + needed = (target_ctas + base_ctas - 1) // base_ctas + sk = 1 + while sk < needed: + sk *= 2 + min_toks_per_part = 64 + max_sk = max(1, KV_MAX // min_toks_per_part) + sk = min(sk, max_sk, 64) + return max(1, sk) + + +def auto_split_k_coop(B: int, G: int, H_q: int, KV_MAX: int) -> int: + """split_k for the gfx950 coop-DMA kernel: latency-bound, wants ~8 waves, cap 64.""" + n_cus = _get_cu_count() + target_ctas = n_cus * 8 # 8 waves + base_ctas = B * G * H_q + needed = max(1, (target_ctas + base_ctas - 1) // base_ctas) + sk = 1 + while sk < needed: + sk *= 2 + MIN_CHUNK_TOKENS = 64 + max_sk = max(1, KV_MAX // MIN_CHUNK_TOKENS) + sk = min(sk, max_sk, 64) + return max(1, sk) + + +def auto_split_k_hp(B: int, G: int, H_q: int, H_kv: int, KV_MAX: int) -> int: + """split_k for the head-packed gfx950 kernel: ~8 waves counted in warps (B*G*H_kv), cap 64.""" + n_cus = _get_cu_count() + target_warps = n_cus * 8 + base_warps = B * G * H_kv + needed = max(1, (target_warps + base_warps - 1) // base_warps) + sk = 1 + while sk < needed: + sk *= 2 + MIN_CHUNK_TOKENS = 64 + max_sk = max(1, KV_MAX // MIN_CHUNK_TOKENS) + sk = min(sk, max_sk, 64) + return max(1, sk) + + +# ── Host launcher ───────────────────────────────────────────────────────────── + + +def pa_decode_launch( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + seq_positions: Optional[torch.Tensor], + softmax_scale: float, + split_k: int = 0, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + """Paged-attention decode (public entry point). Dispatches to gfx950 (head-packed, + ratio 1..16) or gfx950_coop, both falling back to generic off-gfx950.""" + _, _, _, H_q, _ = Q.shape + H_kv = K.shape[3] + ratio = H_q // H_kv if H_kv > 0 else 0 + use_hp = H_kv > 0 and H_q % H_kv == 0 and 1 <= ratio <= 16 + if use_hp: + from .pa_decode_gfx950 import pa_decode_gfx950_launch + + return pa_decode_gfx950_launch( + Q, K, V, seq_positions, softmax_scale, split_k, output_dtype + ) + from .pa_decode_gfx950_coop import pa_decode_gfx950_coop_launch + + return pa_decode_gfx950_coop_launch( + Q, K, V, seq_positions, softmax_scale, split_k, output_dtype + ) + + +# ── AOT interface ───────────────────────────────────────────────────────────── + + +AOT_ARCHS: List[str] = ["gfx942", "gfx950"] + +# KV is f16/bf16 only; split-K path writes f32 partials, so out="f32" for sk>1. +_HEAD_SIZES = (64, 128, 256) +_KV_DTYPES = ("f16", "bf16") +_SPLIT_KS = (1, 2, 4, 8, 16, 32, 64) + +AOT_CONFIGS: List[Dict[str, Any]] = [ + { + "head_size": hs, + "kv_dtype_str": kv, + "output_dtype_str": ("f32" if sk > 1 else kv), + "split_k": sk, + } + for hs in _HEAD_SIZES + for kv in _KV_DTYPES + for sk in _SPLIT_KS +] + + +def compile_aot_config(config: Dict[str, Any], arch: str) -> None: + """Precompile one config. generic on every arch; gfx950 + coop only on gfx950.""" + from .pa_decode_generic import compile_pa_decode_generic + + hs = config["head_size"] + kv = config["kv_dtype_str"] + od = config["output_dtype_str"] + sk = config["split_k"] + + compile_pa_decode_generic( + head_size=hs, + kv_dtype_str=kv, + output_dtype_str=od, + split_k=sk, + arch=arch, + ) + + if arch.startswith("gfx950"): + from .pa_decode_gfx950 import compile_pa_decode_gfx950 + from .pa_decode_gfx950_coop import compile_pa_decode_gfx950_coop + + compile_pa_decode_gfx950_coop( + head_size=hs, + kv_dtype_str=kv, + output_dtype_str=od, + split_k=sk, + arch=arch, + ) + compile_pa_decode_gfx950( + head_size=hs, + kv_dtype_str=kv, + output_dtype_str=od, + split_k=sk, + arch=arch, + ) diff --git a/mslk/attention/fmha/flydsl/pa_decode_fp8.py b/mslk/attention/fmha/flydsl/pa_decode_fp8.py new file mode 100644 index 00000000..822594b8 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_fp8.py @@ -0,0 +1,2448 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL native-FP8 symmetric-scale paged-attention decode (persistent scheduling). + +Ports only the persistent-scheduling small-block compute path (compile_pa_decode_ps ++ pa_decode_ps_kernel) from the upstream FlyDSL reference. + +Grid = (batch, kv_heads, max_context_partition_num); each CTA walks 256-token +sub-partitions with online-softmax loop-carried state. K/V pages come from a +per-sequence block_tables (page sizes 16/64). Query bf16/f16 with kernel-internal +symmetric FP8 query-scale. +""" + +from __future__ import annotations + +import functools +import math +from typing import Any, Dict, List + +import flydsl.compiler as flyc # pyre-ignore[21] +import flydsl.expr as fx # pyre-ignore[21] +import torch +from flydsl._mlir import ir # pyre-ignore[21] +from flydsl._mlir.dialects import llvm # pyre-ignore[21] +from flydsl.compiler.kernel_function import CompilationContext # pyre-ignore[21] +from flydsl.expr import ( # pyre-ignore[21] + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import Int32, T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch as get_hip_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr # pyre-ignore[21] + +from .utils import dpp_xor_f32, maxnumf as _maxnumf, rcp_f32 as _rcp_f32, WARP_SIZE + +# ── Kernel geometry constants ──────────────────────────────────────── +KV_BLOCK_SIZE = 1024 # physical page size (matches SP3 kBlockSize) +KV_COMPUTE_BLOCK = 256 # tile size (matches SP3 kTileKV) +NUM_WARPS = 4 +BLOCK_THREADS = NUM_WARPS * WARP_SIZE # 256 +MFMA_N = 16 +MFMA_K = 32 + +TOKENS_PER_WARP = KV_COMPUTE_BLOCK // NUM_WARPS # 64 +TLOOP = TOKENS_PER_WARP // MFMA_N # 4 +ROWS_PER_WARP = WARP_SIZE // MFMA_N # 4 +FP8_ELEMS_16B = 16 # 16 FP8 per 16-byte load +QKHE_PER_FETCH = FP8_ELEMS_16B * ROWS_PER_WARP # 64 + +VTLOOP = NUM_WARPS # 4 +Q_ELEMS_PER_LANE = 8 +Q_CHUNKS_PER_LANE = Q_ELEMS_PER_LANE // 4 + +# LDS sizes +PROB_ROW_STRIDE_BYTES = 40 # 32 data + 8 padding -> 0 bank conflict +LDS_LOGITS_BYTES = NUM_WARPS * 4 * MFMA_N * PROB_ROW_STRIDE_BYTES # 10240 +LDS_SOFTMAX_BYTES = 2 * NUM_WARPS * MFMA_N * 4 # 512 +LDS_SCALE_V_PADDING = 4 # break K/V same-bank paired writes +LDS_SCALE_V_OFFSET = KV_COMPUTE_BLOCK + LDS_SCALE_V_PADDING +LDS_SCALE_BYTES = ( + LDS_SCALE_V_OFFSET + KV_COMPUTE_BLOCK +) * 4 # K/V per-token scale staging + +FP8_MAX = 240.0 +LOG2E = 1.4426950408889634 + +# Match the Gluon PA decode kernel's AGPR allocation: +# .amdhsa_accum_offset 200, .amdhsa_next_free_vgpr 248 => 48 AGPRs, +# with FP8 MFMA using up to a[44:47]. +PA_MFMA_AGPR_ALLOC = "48,48" +PA_MFMA_AGPR_LLVM_OPTIONS = {"amdgpu-mfma-vgpr-form": False} + +# Tiles per block (1024 tokens / 256 tokens per tile = 4, matches SP3 kNumBlockTiles) +TILES_PER_BLOCK = KV_BLOCK_SIZE // KV_COMPUTE_BLOCK # 4 + +_PACKED_FP8_QUERY_DTYPES = tuple( + dtype + for dtype in ( + torch.uint8, + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e4m3fn", None), + ) + if dtype is not None +) + + +def _cdiv(numer: int, denom: int) -> int: + return (numer + denom - 1) // denom + + +def _pow2_shift(value: int) -> int: + assert value > 0 and (value & (value - 1)) == 0 + return value.bit_length() - 1 + + +def _is_pow2(value: int) -> bool: + return value > 0 and (value & (value - 1)) == 0 + + +def _udiv_pow2(value, divisor: int): + return value >> fx.Int32(_pow2_shift(divisor)) + + +def _urem_pow2(value, divisor: int): + return value & fx.Int32(divisor - 1) + + +def _udiv_const(value, divisor: int): + if const_expr(_is_pow2(divisor)): + return _udiv_pow2(value, divisor) + return value // fx.Int32(divisor) + + +def _urem_const(value, divisor: int): + if const_expr(_is_pow2(divisor)): + return _urem_pow2(value, divisor) + return value % fx.Int32(divisor) + + +def _compute_block_base_dw_i64(phys_block, block_stride, head_offset): + phys_block_i64 = fx.Int64(phys_block) + block_stride_i64 = fx.Int64(block_stride) + head_offset_i64 = fx.Int64(head_offset) + return (phys_block_i64 * block_stride_i64 + head_offset_i64) >> fx.Int64(2) + + +def _extract_global_ptr(tensor): + from flydsl._mlir.dialects import fly as _fly + + raw = ( + tensor.ir_value() + if hasattr(tensor, "ir_value") and not isinstance(tensor, ir.Value) + else tensor + ) + ptr_type = ir.Type.parse("!llvm.ptr<1>") + return _fly.extract_aligned_pointer_as_index(ptr_type, raw) + + +def _global_load_i64x2(global_ptr, byte_offset_i64): + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset_i64), elem_type=T.i8 + ) + return llvm.LoadOp(T.i64x2, ptr, alignment=16).result + + +def _global_load_i32(global_ptr, elem_offset_i32): + byte_offset_i64 = fx.Int64(elem_offset_i32) * fx.Int64(4) + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=byte_offset_i64, elem_type=T.i8 + ) + return llvm.LoadOp(T.i32, ptr, alignment=4).result + + +def _exp2_amdgcn_scalar(scalar_value): + """Direct ``llvm.amdgcn.exp2.f32`` intrinsic (single ``v_exp_f32``) on one + f32 scalar, vs OCML's ``v_exp_f32 + v_ldexp_f32``. Skipping ldexp is safe: + softmax inputs are pre-clamped (safe_qk_max/safe_partition_max) to fast-range. + """ + from flydsl._mlir.ir import F32Type + + raw = ( + arith.unwrap(scalar_value) + if hasattr(scalar_value, "ir_value") or hasattr(scalar_value, "type") + else scalar_value + ) + f32_ty = F32Type.get() + return llvm.call_intrinsic(f32_ty, "llvm.amdgcn.exp2.f32", [raw], [], []) + + +def _exp2_f32_fast(value): + """2^value (elementwise), via the amdgcn intrinsic (see _exp2_amdgcn_scalar).""" + from flydsl._mlir.dialects import vector as _vector_dialect + from flydsl._mlir.ir import VectorType + + raw = ( + arith.unwrap(value) + if hasattr(value, "ir_value") or hasattr(value, "type") + else value + ) + ty = raw.type + if isinstance(ty, VectorType): + n = ty.shape[0] + elems = [] + for i in range(n): + scalar = _vector_dialect.extract( + raw, static_position=[i], dynamic_position=[] + ) + elems.append(_exp2_amdgcn_scalar(scalar)) + return _vector_dialect.from_elements(ty, elems) + return _exp2_amdgcn_scalar(raw) + + +def _unflatten_k(k_flat, qkhe_loop: int = 2): + return [ + [k_flat[td * (qkhe_loop * 2) + j] for j in range(qkhe_loop * 2)] + for td in range(TLOOP) + ] + + +def _flatten_v_results(v_results, vhe_loop: int = 2): + """v_results[vt][vhe] = i64x2 → flat list of scalar i64 (order matches + ``_unflatten_v_results``). Carries V through scf.for state (scalars only).""" + flat = [] + for vt in range(VTLOOP): + for vhe in range(vhe_loop): + v_i64x2 = fx.Vector(v_results[vt][vhe]) + flat.append(v_i64x2[0]) + flat.append(v_i64x2[1]) + return flat + + +def _unflatten_v_results(v_flat, vhe_loop: int = 2): + """Inverse of ``_flatten_v_results``: rebuild v_results[vt][vhe] = i64x2.""" + v_results = [] + idx = 0 + for vt in range(VTLOOP): + vhe_data = [] + for vhe in range(vhe_loop): + v_i64x2 = vector.from_elements( + T.vec(2, T.i64), [v_flat[idx], v_flat[idx + 1]] + ) + vhe_data.append(v_i64x2) + idx += 2 + v_results.append(vhe_data) + return v_results + + +def _build_pa_thread_invariants( + warp_id, + lane16id, + rowid, + *, + trans_v, + per_token_kv, + qkhe_loop: int = 2, + vhe_loop: int = 2, +): + c_tokens_per_warp = fx.Int32(TOKENS_PER_WARP) + c_mfma_n = fx.Int32(MFMA_N) + k_tok_thread_base = warp_id * c_tokens_per_warp + lane16id + c_tok_stride_dw = fx.Int32(FP8_ELEMS_16B // 4) + c_he_stride_dw = fx.Int32(KV_BLOCK_SIZE * FP8_ELEMS_16B // 4) + k_he_off_dw = [ + rowid * c_he_stride_dw + fx.Int32(qkhe * 4) * c_he_stride_dw + for qkhe in range(qkhe_loop) + ] + + vhead_elems = [ + fx.Int32(vhe * NUM_WARPS * MFMA_N) + warp_id * c_mfma_n + lane16id + for vhe in range(vhe_loop) + ] + v_tok_thread_off = [ + fx.Int32(vt * TOKENS_PER_WARP) + rowid * c_mfma_n for vt in range(VTLOOP) + ] + if const_expr(trans_v): + vhead_elem_dw = [ + vhead_elems[vhe] * fx.Int32(FP8_ELEMS_16B // 4) for vhe in range(vhe_loop) + ] + else: + vhead_elem_dw = [ + vhead_elems[vhe] * fx.Int32(KV_BLOCK_SIZE // 4) for vhe in range(vhe_loop) + ] + + kv_tok_thread_base = warp_id * c_tokens_per_warp + rowid * 4 + rowid_8x8 = rowid >> fx.Int32(1) + offset_in_slot = rowid & fx.Int32(1) + prob_wr_thread_base = ( + warp_id * fx.Int32(4 * MFMA_N * PROB_ROW_STRIDE_BYTES) + + lane16id * fx.Int32(PROB_ROW_STRIDE_BYTES) + + rowid_8x8 * fx.Int32(8) + + offset_in_slot * 4 + ) + pv_prob_read_base = rowid * fx.Int32( + MFMA_N * PROB_ROW_STRIDE_BYTES + ) + lane16id * fx.Int32(PROB_ROW_STRIDE_BYTES) + + sm_lane_wave_base = lane16id * fx.Int32(NUM_WARPS) + sm_max_off = fx.Index(sm_lane_wave_base + warp_id) + sm_sum_off = fx.Index(fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id) + sm_rd_max_offs = [ + fx.Index(sm_lane_wave_base + fx.Int32(w)) for w in range(NUM_WARPS) + ] + sm_rd_sum_offs = [ + fx.Index(fx.Int32(NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w)) + for w in range(NUM_WARPS) + ] + + sm_vmax_wr_off = None + sm_vmax_rd_offs = None + if const_expr(per_token_kv): + sm_vmax_wr_off = fx.Index( + fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + warp_id + ) + sm_vmax_rd_offs = [ + fx.Index(fx.Int32(2 * NUM_WARPS * MFMA_N) + sm_lane_wave_base + fx.Int32(w)) + for w in range(NUM_WARPS) + ] + + return ( + k_tok_thread_base, + c_tok_stride_dw, + k_he_off_dw, + v_tok_thread_off, + vhead_elem_dw, + kv_tok_thread_base, + prob_wr_thread_base, + pv_prob_read_base, + sm_max_off, + sm_sum_off, + sm_rd_max_offs, + sm_rd_sum_offs, + sm_vmax_wr_off, + sm_vmax_rd_offs, + ) + + +def _compute_mtp_group_state( + lane16id, + local_qhead_idx, + *, + mtp_group_idx, + query_length, + query_group_size, +): + g_off = mtp_group_idx * 16 + lane_pair_raw = lane16id + fx.Int32(g_off) + c_total_pairs = fx.Int32(query_length * query_group_size) + c_pair_max = fx.Int32(query_length * query_group_size - 1) + c_ql_m1 = fx.Int32(query_length - 1) + + if const_expr((query_length * query_group_size) % MFMA_N == 0): + lane_pair = lane_pair_raw + else: + lane_pair = arith.select( + lane_pair_raw < c_total_pairs, lane_pair_raw, c_pair_max + ) + qi_raw = _udiv_const(lane_pair, query_group_size) + if const_expr((query_length * query_group_size) % MFMA_N == 0): + qi_val = qi_raw + else: + qi_val = arith.select(qi_raw < c_ql_m1, qi_raw, c_ql_m1) + qhi_pos = _urem_const(lane_pair, query_group_size) + + lqh_pair_raw = local_qhead_idx + fx.Int32(g_off) + if const_expr((query_length * query_group_size) % MFMA_N == 0): + lqh_pair = lqh_pair_raw + else: + lqh_pair = arith.select(lqh_pair_raw < c_total_pairs, lqh_pair_raw, c_pair_max) + lqi_raw = _udiv_const(lqh_pair, query_group_size) + if const_expr((query_length * query_group_size) % MFMA_N == 0): + qi_for_q = lqi_raw + else: + qi_for_q = arith.select(lqi_raw < c_ql_m1, lqi_raw, c_ql_m1) + local_qhead_idx_for_q = _urem_const(lqh_pair, query_group_size) + return qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q + + +@flyc.jit +def _prefetch_q_chunks( + q_rsrc, + q_base, + lane16id, + *, + query_load_is_bf16, + q_lanes_per_head, + q_elems_per_lane: int = Q_ELEMS_PER_LANE, + q_chunks_per_lane: int = Q_CHUNKS_PER_LANE, +): + # Each lane owns q_elems_per_lane (= max(8, head_dim//MFMA_N)) contiguous Q elems + # via q_chunks_per_lane vec4 loads; 16 lanes cover the full head-dim. + q_load_lane = lane16id + if const_expr(q_lanes_per_head < MFMA_N): + q_load_lane = arith.select( + lane16id < fx.Int32(q_lanes_per_head), lane16id, fx.Int32(0) + ) + q_elem = q_base + q_load_lane * fx.Int32(q_elems_per_lane) + q_chunks = [] + for qwi in range_constexpr(q_chunks_per_lane): + q_chunks.append( + buffer_ops.buffer_load( + q_rsrc, + q_elem + fx.Int32(qwi * 4), + vec_width=4, + dtype=fx.BFloat16 if query_load_is_bf16 else fx.Float16, + ) + ) + return q_chunks + + +@flyc.jit +def _finish_q_fragments( + logits_lds_i32, + logits_lds_i64, + softmax_lds_f32, + q_chunks, + lane16id, + rowid, + local_qhead_idx, + *, + head_size: int, + qkhe_loop: int, + q_lanes_per_head: int, + q_elems_per_lane: int = Q_ELEMS_PER_LANE, +): + # LDS Q layout (per-qhead contiguous): Q[head=h][hd=d] at byte offset + # h*HEAD_SIZE + d (FP8), aliased with later P writes via logits_lds_*. + # Writer: qhead=local_qhead_idx writes 1 i64 at qhead*HEAD_SIZE + lane16id*8. + # Reader (mfma_f32_16x16x32_fp8_fp8, B=Q^T): thread (rowid R, lane L), + # k_step=qkhe*2+qkr, byte offset L*HEAD_SIZE + qkhe*64 + R*16 + qkr*8. + c_head_size = fx.Int32(head_size) + lds_q_base = local_qhead_idx * c_head_size + lane16id * fx.Int32(q_elems_per_lane) + abs_mask = fx.Vector.filled(4, 0x7FFFFFFF, fx.Int32) + c_zero_f = fx.Float32(0.0) + c_one_f = fx.Float32(1.0) + + q_f32_chunks = [] + local_max = c_zero_f + for q_src in q_chunks: + q_f32 = fx.Vector(q_src).to(fx.Float32) + q_f32_chunks.append(q_f32) + q_i32 = q_f32.bitcast(fx.Int32) + q_abs_i32 = q_i32 & abs_mask + q_abs = q_abs_i32.bitcast(fx.Float32) + chunk_max = q_abs.reduce("max") + local_max = _maxnumf(local_max, chunk_max) + + for sh in [8, 4, 2, 1]: + local_max = _maxnumf(local_max, dpp_xor_f32(local_max, sh)) + query_scale_lane = fx.Float32( + arith.select( + local_max > c_zero_f, + local_max * fx.Float32(1.0 / FP8_MAX).ir_value(), + c_one_f, + ) + ) + inv_query_scale = _rcp_f32(query_scale_lane) + q_words = [] + for q_f32 in q_f32_chunks: + p = q_f32 * inv_query_scale + lo = rocdl.cvt_pk_fp8_f32(T.i32, p[0], p[1], fx.Int32(0), False) + q_words.append(rocdl.cvt_pk_fp8_f32(T.i32, p[2], p[3], lo, True)) + if lane16id == fx.Int32(0): + fx.Vector.from_elements([query_scale_lane], dtype=fx.Float32).store( + softmax_lds_f32, [fx.Index(local_qhead_idx)] + ) + + # One packed i32 word per vec4 chunk, stored as a single vec at the per-lane base. + v01 = fx.Vector.from_elements(q_words, dtype=fx.Int32) + lds_q_i32 = lds_q_base >> fx.Int32(2) + if const_expr(q_lanes_per_head < MFMA_N): + if lane16id < fx.Int32(q_lanes_per_head): + v01.store(logits_lds_i32, [fx.Index(lds_q_i32)]) + else: + v01.store(logits_lds_i32, [fx.Index(lds_q_i32)]) + + q_frags = [] + gpu.barrier() + query_scale_lane = fx.Vector.load( + T.vec(1, fx.Float32.ir_type), softmax_lds_f32, [fx.Index(lane16id)] + )[0].ir_value() + for qkhe in range_constexpr(qkhe_loop): + for qkr in range_constexpr(2): + # Byte offset: lane16id*HEAD_SIZE + qkhe*64 + rowid*16 + qkr*8 (see above). + lds_rd_byte = ( + lane16id * c_head_size + + fx.Int32(qkhe << 6) + + (rowid << fx.Int32(4)) + + fx.Int32(qkr << 3) + ) + lds_rd_base = lds_rd_byte >> fx.Int32(3) + q_v1 = fx.Vector.load( + T.vec(1, T.i64), logits_lds_i64, [fx.Index(lds_rd_base)] + ) + q_frags.append(q_v1[0]) + return q_frags, query_scale_lane + + +def _prefetch_mtp_group_query( + q_rsrc, + batch_idx, + kv_h, + stride_q_seq, + stride_q_head, + lane16id, + local_qhead_idx, + *, + mtp_group_idx, + query_length, + query_group_size, + query_load_is_bf16, + q_lanes_per_head, + q_elems_per_lane: int = Q_ELEMS_PER_LANE, + q_chunks_per_lane: int = Q_CHUNKS_PER_LANE, +): + qi_val, qhi_pos, qi_for_q, local_qhead_idx_for_q = _compute_mtp_group_state( + lane16id, + local_qhead_idx, + mtp_group_idx=mtp_group_idx, + query_length=query_length, + query_group_size=query_group_size, + ) + q_row = batch_idx * arith.constant(query_length, type=T.i32) + qi_for_q + q_base = ( + q_row * stride_q_seq + + (kv_h * arith.constant(query_group_size, type=T.i32) + local_qhead_idx_for_q) + * stride_q_head + ) + q_chunks = _prefetch_q_chunks( + q_rsrc, + q_base, + lane16id, + query_load_is_bf16=query_load_is_bf16, + q_lanes_per_head=q_lanes_per_head, + q_elems_per_lane=q_elems_per_lane, + q_chunks_per_lane=q_chunks_per_lane, + ) + return qi_val, qhi_pos, q_chunks + + +def _finish_mtp_group_q_fragments( + logits_lds_i32, + logits_lds_i64, + softmax_lds_f32, + mtp_prefetch, + lane16id, + rowid, + local_qhead_idx, + *, + head_size: int, + qkhe_loop: int, + q_lanes_per_head: int, + q_elems_per_lane: int = Q_ELEMS_PER_LANE, +): + qi_val, qhi_pos, q_chunks = mtp_prefetch + q_frags, query_scale_lane = _finish_q_fragments( + logits_lds_i32, + logits_lds_i64, + softmax_lds_f32, + q_chunks, + lane16id, + rowid, + local_qhead_idx, + head_size=head_size, + qkhe_loop=qkhe_loop, + q_lanes_per_head=q_lanes_per_head, + q_elems_per_lane=q_elems_per_lane, + ) + return qi_val, qhi_pos, q_frags, query_scale_lane + + +def _normalize_pa_output(running_sum, outs, zero_f): + one_f = fx.Float32(1.0).ir_value() + safe_sum = arith.select(running_sum > zero_f, running_sum, one_f) + inv_sum = _rcp_f32(safe_sum) + inv_sum_vec = vector.broadcast(T.f32x4, inv_sum) + return [out * inv_sum_vec for out in outs] + + +@flyc.jit +def _make_pa_phase_helpers( + *, + trans_v, + per_token_q, + per_token_kv, + needs_mask, + query_length, + kv_h, + v_global_ptr, + ks_rsrc, + vs_rsrc, + logits_lds_i32, + logits_lds_i64, + softmax_lds_f32, + scale_lds_f32, + stride_ks_block, + stride_ks_head, + softmax_scale_base, + softmax_q_scale, + k_scale_val, + scale, + v_scale_val, + warp_id, + lane16id, + rowid, + k_tok_thread_base, + v_tok_thread_off, + vhead_elem_dw, + kv_tok_thread_base, + prob_wr_thread_base, + pv_prob_read_base, + sm_max_off, + sm_sum_off, + sm_rd_max_offs, + sm_rd_sum_offs, + sm_vmax_wr_off, + sm_vmax_rd_offs, + c_w, + neg_inf, + zero_f, + cache_scale_vecs=False, + head_size: int = 128, + qkhe_loop: int = 2, + vhe_loop: int = 2, +): + apply_causal_mask = needs_mask or query_length > 1 + pv_prob_i64_indices = [] + for vt in range_constexpr(VTLOOP): + for j in range_constexpr(2): + p_byte = ( + arith.constant(vt * 4 * MFMA_N * PROB_ROW_STRIDE_BYTES, type=T.i32) + + pv_prob_read_base + + arith.constant(j * 8, type=T.i32) + ) + pv_prob_i64_indices.append(fx.Index(p_byte >> fx.Int32(3))) + + def _load_kv_scale_scalars(tile_token_offset_i32, phys_block): + if const_expr(per_token_kv): + scale_block_base = phys_block * stride_ks_block + kv_h * stride_ks_head + scale_stage_token = ( + warp_id * fx.Int32(WARP_SIZE) + rowid * fx.Int32(MFMA_N) + lane16id + ) + scale_global_token = tile_token_offset_i32 + scale_stage_token + k_scale_scalar = buffer_ops.buffer_load( + ks_rsrc, + scale_block_base + scale_global_token, + vec_width=1, + dtype=fx.Float32, + ) + v_scale_scalar = buffer_ops.buffer_load( + vs_rsrc, + scale_block_base + scale_global_token, + vec_width=1, + dtype=fx.Float32, + ) + return k_scale_scalar, v_scale_scalar + return None + + def _load_v_and_scales( + v_block_base_dw, + tile_token_offset_i32, + *, + phys_block, + preloaded_scale_scalars=None, + ): + if const_expr(per_token_kv): + scale_stage_token = ( + warp_id * fx.Int32(WARP_SIZE) + rowid * fx.Int32(MFMA_N) + lane16id + ) + if const_expr(preloaded_scale_scalars is None): + preloaded_scale_scalars = _load_kv_scale_scalars( + tile_token_offset_i32, phys_block + ) + k_scale_scalar, v_scale_scalar = preloaded_scale_scalars + fx.Vector.from_elements([k_scale_scalar], dtype=fx.Float32).store( + scale_lds_f32, + [fx.Index(scale_stage_token)], + ) + fx.Vector.from_elements([v_scale_scalar], dtype=fx.Float32).store( + scale_lds_f32, + [fx.Index(fx.Int32(LDS_SCALE_V_OFFSET) + scale_stage_token)], + ) + rocdl.sched_barrier(0) + + v_results = [] + for vt in range_constexpr(VTLOOP): + vhe_data = [] + for vhe in range_constexpr(vhe_loop): + v_token_in_block = tile_token_offset_i32 + v_tok_thread_off[vt] + if const_expr(trans_v): + vt_group = v_token_in_block >> fx.Int32(4) + va_dw_delta = ( + vt_group + * arith.constant(head_size * FP8_ELEMS_16B // 4, type=T.i32) + + vhead_elem_dw[vhe] + ) + else: + va_dw_delta = vhead_elem_dw[vhe] + (v_token_in_block >> fx.Int32(2)) + va_byte = (v_block_base_dw + fx.Int64(va_dw_delta)) * fx.Int64(4) + v_i64x2 = _global_load_i64x2(v_global_ptr, va_byte) + vhe_data.append(v_i64x2) + v_results.append(vhe_data) + + if const_expr(per_token_kv): + gpu.barrier() + if const_expr(cache_scale_vecs): + k_scale_vecs = [] + v_scale_vecs = [] + for td in range_constexpr(TLOOP): + scale_row_base = kv_tok_thread_base + fx.Int32(td * MFMA_N) + k_scale_vecs.append( + vector.load_op( + T.f32x4, scale_lds_f32, [fx.Index(scale_row_base)] + ) + ) + v_scale_vecs.append( + vector.load_op( + T.f32x4, + scale_lds_f32, + [fx.Index(fx.Int32(LDS_SCALE_V_OFFSET) + scale_row_base)], + ) + ) + return v_results, k_scale_vecs, v_scale_vecs + + return v_results + + def _scale_row_base(td: int): + return kv_tok_thread_base + fx.Int32(td * MFMA_N) + + def _load_k_scale_vec(td: int): + return vector.load_op(T.f32x4, scale_lds_f32, [fx.Index(_scale_row_base(td))]) + + def _load_v_scale_vec(td: int): + return vector.load_op( + T.f32x4, + scale_lds_f32, + [fx.Index(fx.Int32(LDS_SCALE_V_OFFSET) + _scale_row_base(td))], + ) + + def _get_k_scale_vec(td: int, k_scale_vecs=None): + if const_expr(cache_scale_vecs): + return k_scale_vecs[td] + return _load_k_scale_vec(td) + + def _get_v_scale_vec(td: int, v_scale_vecs=None): + if const_expr(cache_scale_vecs): + return v_scale_vecs[td] + return _load_v_scale_vec(td) + + def _store_vmax_warp(partition_start, *, seq_end=None, v_scale_vecs=None): + if const_expr(per_token_kv): + kv_tok_base = ( + partition_start + kv_tok_thread_base + if const_expr(seq_end is not None) + else None + ) + v_max_warp = zero_f + for td in range_constexpr(TLOOP): + vs = _get_v_scale_vec(td, v_scale_vecs) + for i in range_constexpr(4): + if const_expr(kv_tok_base is not None): + kv_tok = kv_tok_base + arith.constant( + td * MFMA_N + i, type=T.i32 + ) + vs_i = vector.extract( + vs, static_position=[i], dynamic_position=[] + ) + vs_i = arith.select(kv_tok < seq_end, vs_i, zero_f) + vs = vector.insert( + vs_i, vs, static_position=[i], dynamic_position=[] + ) + v_max_warp = _maxnumf(v_max_warp, fx.Vector(vs).reduce("max")) + for sh in [32, 16]: + v_max_warp = _maxnumf( + v_max_warp, + v_max_warp.shuffle_xor(arith.constant(sh, type=T.i32), c_w), + ) + vector.store( + fx.Vector.from_elements([v_max_warp], dtype=fx.Float32), + softmax_lds_f32, + [sm_vmax_wr_off], + ) + + def _token_vec_i32(kv_tok_base, td: int): + kv_tok_td_base = kv_tok_base + arith.constant(td * MFMA_N, type=T.i32) + return fx.Vector.from_elements( + [ + kv_tok_td_base + arith.constant(i, type=T.i32) + for i in range_constexpr(4) + ], + dtype=fx.Int32, + ) + + def _apply_token_mask_vec( + logit_vec, td: int, kv_tok_base, causal_bound, false_value + ): + tok_vec = _token_vec_i32(kv_tok_base, td) + if const_expr(apply_causal_mask): + in_range = tok_vec < causal_bound + return arith.select( + in_range, + logit_vec, + vector.broadcast(T.f32x4, arith.unwrap(false_value)), + ) + return logit_vec + + def _qk_and_intra_softmax( + k_ops, + partition_start, + q_frags, + causal_bound, + query_scale_lane=None, + *, + preloaded_scales=None, + ): + if const_expr(preloaded_scales is not None): + if const_expr(cache_scale_vecs and per_token_kv): + k_scale_vecs, v_scale_vecs = preloaded_scales + + query_scale_vec = None + if const_expr(per_token_q): + query_scale_vec = vector.broadcast( + T.f32x4, query_scale_lane * softmax_scale_base + ) + d_out = [] + for td in range_constexpr(TLOOP): + acc = arith.constant_vector(0.0, T.f32x4) + for k_step in range_constexpr(qkhe_loop * 2): + acc = rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, [k_ops[td][k_step], q_frags[k_step], acc, 0, 0, 0] + ) + if const_expr(per_token_kv): + if const_expr(cache_scale_vecs and per_token_kv): + k_scale_vec = _get_k_scale_vec(td, k_scale_vecs) + else: + k_scale_vec = _get_k_scale_vec(td) + scale_vec = ( + k_scale_vec * query_scale_vec + if const_expr(per_token_q) + else k_scale_vec * vector.broadcast(T.f32x4, softmax_q_scale) + ) + d_out.append(acc * scale_vec) + else: + if const_expr(per_token_q): + d_out.append( + acc * (query_scale_vec * vector.broadcast(T.f32x4, k_scale_val)) + ) + else: + d_out.append(acc * vector.broadcast(T.f32x4, scale)) + + kv_tok_base = ( + partition_start + kv_tok_thread_base + if const_expr(apply_causal_mask) + else None + ) + qk_max = neg_inf + for td in range_constexpr(TLOOP): + logits_vec = d_out[td] + if const_expr(kv_tok_base is not None): + logits_vec = _apply_token_mask_vec( + logits_vec, td, kv_tok_base, causal_bound, neg_inf + ) + d_out[td] = logits_vec + qk_max = _maxnumf(qk_max, fx.Vector(logits_vec).reduce("max")) + for sh in [32, 16]: + qk_max = _maxnumf( + qk_max, qk_max.shuffle_xor(arith.constant(sh, type=T.i32), c_w) + ) + vector.store( + fx.Vector.from_elements([qk_max], dtype=fx.Float32), + softmax_lds_f32, + [sm_max_off], + ) + + if const_expr(cache_scale_vecs and per_token_kv): + return d_out, v_scale_vecs + return d_out + + def _cross_warp_softmax_and_prob_pack(d_out, rmax, rsum, outs, v_scale_vecs): + partition_max = neg_inf + partition_sum = zero_f + max_vec = fx.Vector( + vector.load_op(T.f32x4, softmax_lds_f32, [sm_rd_max_offs[0]]) + ) + for w in range_constexpr(NUM_WARPS): + partition_max = _maxnumf(partition_max, max_vec[w]) + + new_rmax = _maxnumf(rmax, partition_max) + safe_eff_max = ( + arith.select(partition_max > neg_inf, new_rmax, zero_f) + if const_expr(needs_mask) + else new_rmax + ) + local_exp_sum = zero_f + for td in range_constexpr(TLOOP): + diff_vec = fx.Vector(d_out[td]) - vector.broadcast( + T.f32x4, arith.unwrap(safe_eff_max) + ) + p_vec = _exp2_f32_fast( + diff_vec * vector.broadcast(T.f32x4, arith.unwrap(fx.Float32(LOG2E))) + ) + local_exp_sum = local_exp_sum + fx.Vector(p_vec).reduce("add") + d_out[td] = p_vec + for sh in [32, 16]: + local_exp_sum = local_exp_sum + local_exp_sum.shuffle_xor( + arith.constant(sh, type=T.i32), c_w + ) + vector.store( + fx.Vector.from_elements([local_exp_sum], dtype=fx.Float32), + softmax_lds_f32, + [sm_sum_off], + ) + if const_expr(needs_mask): + accum_scale = arith.select( + rmax > neg_inf, + _exp2_f32_fast((rmax - new_rmax) * fx.Float32(LOG2E).ir_value()), + zero_f, + ) + else: + accum_scale = _exp2_f32_fast( + (rmax - new_rmax) * fx.Float32(LOG2E).ir_value() + ) + + gpu.barrier() + sum_vec = fx.Vector( + vector.load_op(T.f32x4, softmax_lds_f32, [sm_rd_sum_offs[0]]) + ) + for w in range_constexpr(NUM_WARPS): + partition_sum = arith.addf( + arith.unwrap(partition_sum), + arith.unwrap(sum_vec[w]), + fastmath=arith.FastMathFlags.contract, + ) + + accum_sum = arith.mulf( + arith.unwrap(accum_scale), + arith.unwrap(rsum), + fastmath=arith.FastMathFlags.contract, + ) + rsum = arith.addf( + accum_sum, + arith.unwrap(partition_sum), + fastmath=arith.FastMathFlags.contract, + ) + rmax = new_rmax + accum_scale_vec = vector.broadcast(T.f32x4, arith.unwrap(accum_scale)) + for vhe in range_constexpr(vhe_loop): + outs[vhe] = outs[vhe] * accum_scale_vec + + if const_expr(per_token_kv): + v_max_global = zero_f + vmax_vec = fx.Vector( + vector.load_op(T.f32x4, softmax_lds_f32, [sm_vmax_rd_offs[0]]) + ) + for w in range_constexpr(NUM_WARPS): + w_vmax = vmax_vec[w] + v_max_global = _maxnumf(v_max_global, w_vmax) + v_max_scaled = v_max_global * fx.Float32(1.0 / FP8_MAX).ir_value() + v_max_safe_scaled = v_max_scaled + fx.Float32(1e-8 / FP8_MAX).ir_value() + norm_factor = _rcp_f32(v_max_safe_scaled) + v_correction = v_max_scaled + _vec_norm_p = arith.unwrap(norm_factor) + for td in range_constexpr(TLOOP): + d_out[td] = d_out[td] * ( + _get_v_scale_vec(td, v_scale_vecs) + * vector.broadcast(T.f32x4, _vec_norm_p) + ) + else: + v_correction = v_scale_val + + for td in range_constexpr(TLOOP): + p0 = vector.extract(d_out[td], static_position=[0], dynamic_position=[]) + p1 = vector.extract(d_out[td], static_position=[1], dynamic_position=[]) + p2 = vector.extract(d_out[td], static_position=[2], dynamic_position=[]) + p3 = vector.extract(d_out[td], static_position=[3], dynamic_position=[]) + lo = rocdl.cvt_pk_fp8_f32( + T.i32, p0, p1, arith.constant(0, type=T.i32), False + ) + pk = rocdl.cvt_pk_fp8_f32(T.i32, p2, p3, lo, True) + byte_base = prob_wr_thread_base + arith.constant( + td * MFMA_N * PROB_ROW_STRIDE_BYTES, type=T.i32 + ) + i32_off = byte_base >> fx.Int32(2) + pk_vec = vector.from_elements(T.vec(1, T.i32), [pk]) + vector.store(pk_vec, logits_lds_i32, [fx.Index(i32_off)]) + return rmax, rsum, outs, v_correction + + def _pv_mfma(v_ops, outs, v_correction): + v_correction = fx.Float32(v_correction).ir_value() + fm_contract = arith.FastMathFlags.contract + v_correction_vec = vector.broadcast(T.f32x4, v_correction) + + # Hoist all P_i64 LDS loads out of the vhe loop (p_i64 depends only on + # (vt, j)) so the compiler can pipeline them before the MFMA chain. + p_i64_all = [] + for vt in range_constexpr(VTLOOP): + for j in range_constexpr(2): + p_i64_idx = pv_prob_i64_indices[vt * 2 + j] + p_i64_all.append( + fx.Vector.load(T.vec(1, T.i64), logits_lds_i64, [p_i64_idx])[0] + ) + + for vhe in range_constexpr(vhe_loop): + tmp_out = arith.constant_vector(0.0, T.f32x4) + for vt in range_constexpr(VTLOOP): + v_i64x2 = fx.Vector(v_ops[vt][vhe]) + for j in range_constexpr(2): + tmp_out = rocdl.mfma_f32_16x16x32_fp8_fp8( + T.f32x4, + [ + v_i64x2[j], + p_i64_all[vt * 2 + j], + tmp_out, + 0, + 0, + 0, + ], + ) + outs[vhe] = arith.addf( + arith.mulf(tmp_out, v_correction_vec, fastmath=fm_contract), + outs[vhe], + fastmath=fm_contract, + ) + return outs + + return ( + _load_kv_scale_scalars, + _load_v_and_scales, + _store_vmax_warp, + _qk_and_intra_softmax, + _cross_warp_softmax_and_prob_pack, + _pv_mfma, + ) + + +def _is_current_stream_capturing() -> bool: + if not torch.cuda.is_available(): + return False + try: + return torch.cuda.is_current_stream_capturing() + except RuntimeError: + return False + + +def _prepare_scale_tensor( + name: str, + scale, + *, + device: torch.device, + is_graph_capturing: bool, +) -> torch.Tensor: + if isinstance(scale, torch.Tensor): + if is_graph_capturing: + if scale.device != device: + raise ValueError( + f"CUDA graph capture requires `{name}` to already be on {device}, " + f"got {scale.device}." + ) + if scale.dtype != torch.float32: + raise ValueError( + f"CUDA graph capture requires `{name}` to already be float32, " + f"got {scale.dtype}." + ) + return scale + return scale.to(device=device, dtype=torch.float32) + + if is_graph_capturing: + raise ValueError( + f"CUDA graph capture requires `{name}` to be passed as a pre-created " + "float32 tensor on the target device." + ) + + return torch.tensor([float(scale or 1.0)], device=device, dtype=torch.float32) + + +def _get_query_input_dtype(query: torch.Tensor) -> str: + if query.dtype in _PACKED_FP8_QUERY_DTYPES: + return "packed_fp8" + if query.dtype == torch.bfloat16: + return "bf16" + if query.dtype == torch.float16: + return "f16" + raise ValueError( + f"Unsupported query dtype for pa_decode_ps_launch: {query.dtype}. " + "Expected packed FP8/uint8, bf16, or f16." + ) + + +def _get_output_dtype_str(output: torch.Tensor) -> str: + if output.dtype == torch.bfloat16: + return "bf16" + if output.dtype == torch.float16: + return "f16" + if output.dtype == torch.float32: + return "f32" + raise ValueError( + f"Unsupported output dtype for pa_decode_ps_launch reduce: {output.dtype}. " + "Expected bf16, f16, or f32." + ) + + +def get_recommended_splits( + num_sequences: int, + num_kv_heads: int, + split_kv_blocks: int = 1, + *, + sliding_window: int = 0, + context_partition_size: int = KV_COMPUTE_BLOCK, + query_length: int = 1, +) -> int: + """Recommend max_context_partition_num; mirrors aiter's get_recommended_splits.""" + if sliding_window > 0: + window_token_count = sliding_window + query_length + return _cdiv(window_token_count - 1, context_partition_size) + 1 + + props = torch.cuda.get_device_properties(torch.device("cuda")) + occupancy = 2 # matches reference Gluon get_occupancy() + num_sm = props.multi_processor_count * occupancy + denom = max(1, num_sequences * num_kv_heads * split_kv_blocks) + n = _cdiv(num_sm, denom) * split_kv_blocks + return max(4, min(n, 8)) + + +# block_size 16/64 handled directly by the small-block PS path here (the +# reference routes them through the metadata worklist path). +_PA_DECODE_PS_SMALL_BLOCK_SIZES = (16, 64) + + +@flyc.jit +def _pa_small_block_load_k_flat( + k_global_ptr, + kv_h_i32, + stride_k_block_i32, + stride_k_head_i32, + lane16id_i32, + rowid_i32, + *, + block_size: int, + phys_blocks, + qkhe_loop: int = 2, +): + """Load K for one warp's 64-token slice of a 256-token partition. Returns + ``k_flat`` (TLOOP * qkhe_loop * 2 i64 scalars) for ``_unflatten_k``/MFMA. + """ + c_he_stride_dw = fx.Int32(block_size * FP8_ELEMS_16B // 4) + c_tok_stride_dw = fx.Int32(FP8_ELEMS_16B // 4) + k_he_off_dw = [ + rowid_i32 * c_he_stride_dw + fx.Int32(qkhe * 4) * c_he_stride_dw + for qkhe in range(qkhe_loop) + ] + k_head_off = kv_h_i32 * stride_k_head_i32 + + k_flat = [] + if const_expr(block_size == 64): + # Each warp owns exactly one physical block (64 tokens). + phys_block = phys_blocks + k_block_base_dw = _compute_block_base_dw_i64( + phys_block, stride_k_block_i32, k_head_off + ) + for td in range_constexpr(TLOOP): + within_block_token = fx.Int32(td * MFMA_N) + lane16id_i32 + kbo_dw = within_block_token * c_tok_stride_dw + for qkhe in range_constexpr(qkhe_loop): + ka_dw = k_block_base_dw + fx.Int64(kbo_dw + k_he_off_dw[qkhe]) + k2 = _global_load_i64x2(k_global_ptr, ka_dw * fx.Int64(4)) + k2_words = fx.Vector(k2) + k_flat.append(k2_words[0]) + k_flat.append(k2_words[1]) + else: + # block_size == 16: each warp spans 4 blocks (one MFMA tile per block). + within_block_token = lane16id_i32 + kbo_dw = within_block_token * c_tok_stride_dw + for td in range_constexpr(TLOOP): + phys_block = phys_blocks[td] + k_block_base_dw = _compute_block_base_dw_i64( + phys_block, stride_k_block_i32, k_head_off + ) + for qkhe in range_constexpr(qkhe_loop): + ka_dw = k_block_base_dw + fx.Int64(kbo_dw + k_he_off_dw[qkhe]) + k2 = _global_load_i64x2(k_global_ptr, ka_dw * fx.Int64(4)) + rocdl.sched_barrier(rocdl.mask_vmem_rd) + k2_words = fx.Vector(k2) + k_flat.append(k2_words[0]) + k_flat.append(k2_words[1]) + return k_flat + + +@flyc.jit +def _pa_small_block_load_v_trans( + v_global_ptr, + kv_h_i32, + stride_v_block_i32, + stride_v_head_i32, + warp_id_i32, + lane16id_i32, + rowid_i32, + v_phys_blocks, + *, + block_size: int, + head_size: int = 128, + vhe_loop: int = 2, +): + """Load V tiles for one CTA's 256-token partition (``trans_v=True``). + Returns ``v_results[vt][vhe]`` (i64x2) indexed as ``_load_v_and_scales``. + """ + v_head_off = kv_h_i32 * stride_v_head_i32 + vhead_elems = [ + fx.Int32(vhe * NUM_WARPS * MFMA_N) + + warp_id_i32 * fx.Int32(MFMA_N) + + lane16id_i32 + for vhe in range(vhe_loop) + ] + vhead_elem_dw = [ + vhead_elems[vhe] * fx.Int32(FP8_ELEMS_16B // 4) for vhe in range(vhe_loop) + ] + c_subblock_dw = fx.Int32(head_size * FP8_ELEMS_16B // 4) + + v_results = [] + for vt in range_constexpr(VTLOOP): + phys_block = v_phys_blocks[vt] + if const_expr(block_size == 64): + # vt selects the physical block (4 blocks per partition); rowid + # selects the 16-token sub-block within that physical block. + sub_block_idx = rowid_i32 + else: + # block_size == 16: (vt * 4 + rowid) selects the block; only one + # 16-token sub-block per physical block, so sub_block_idx == 0. + sub_block_idx = fx.Int32(0) + v_block_base_dw = _compute_block_base_dw_i64( + phys_block, stride_v_block_i32, v_head_off + ) + vhe_data = [] + for vhe in range_constexpr(vhe_loop): + va_dw_delta = sub_block_idx * c_subblock_dw + vhead_elem_dw[vhe] + va_byte = (v_block_base_dw + fx.Int64(va_dw_delta)) * fx.Int64(4) + v_i64x2 = _global_load_i64x2(v_global_ptr, va_byte) + vhe_data.append(v_i64x2) + v_results.append(vhe_data) + return v_results + + +@functools.lru_cache(maxsize=256) +def compile_pa_decode_ps( + *, + block_size: int, + max_context_partition_num: int, + softmax_scale: float = None, + trans_v: bool = True, + query_group_size: int = 16, + per_token_kv: bool = False, + query_length: int = 1, + query_input_dtype: str = "bf16", + head_dim: int = 128, +): + """Compile the small-block partition kernel. See module-level comment.""" + if block_size not in _PA_DECODE_PS_SMALL_BLOCK_SIZES: + raise ValueError( + f"compile_pa_decode_ps: unsupported block_size={block_size}; " + f"expected one of {_PA_DECODE_PS_SMALL_BLOCK_SIZES}." + ) + if query_input_dtype not in ("bf16", "f16"): + raise ValueError( + "compile_pa_decode_ps currently expects bf16/f16 query inputs." + ) + if not trans_v: + raise NotImplementedError( + "compile_pa_decode_ps: trans_v=False not yet supported." + ) + if ( + head_dim % QKHE_PER_FETCH != 0 + or head_dim % (MFMA_N * NUM_WARPS) != 0 + or head_dim % Q_ELEMS_PER_LANE != 0 + ): + raise ValueError( + f"Unsupported head_dim={head_dim}; must be a multiple of {MFMA_N * NUM_WARPS}." + ) + _HEAD = head_dim + _QKHELOOP = head_dim // QKHE_PER_FETCH + _VHELOOP = head_dim // MFMA_N // NUM_WARPS + # Clamp to 8 so head_dim<=128 keeps the fixed 8/2 path. + _Q_ELEMS_PER_LANE = max(Q_ELEMS_PER_LANE, head_dim // MFMA_N) + _Q_CHUNKS_PER_LANE = _Q_ELEMS_PER_LANE // 4 + _Q_LANES_PER_HEAD = head_dim // _Q_ELEMS_PER_LANE + _N_K_h = TLOOP * _QKHELOOP * 2 + _N_V_FLAT_h = 2 * VTLOOP * _VHELOOP + + arch = get_hip_arch() + query_load_is_bf16 = query_input_dtype == "bf16" + if softmax_scale is None: + softmax_scale = 1.0 / (head_dim**0.5) + _softmax_scale = float(softmax_scale) + _block_size = block_size + _blocks_per_partition = KV_COMPUTE_BLOCK // _block_size + + _mtp_groups = max(1, math.ceil(query_length * query_group_size / 16)) + + # LDS allocation. per_token_kv adds a cross-warp v_scale_max region + # (appended to softmax) and a K/V per-token scale staging region. + LDS_VMAX_BYTES = NUM_WARPS * MFMA_N * 4 if const_expr(per_token_kv) else 0 + LDS_SOFTMAX_TOTAL = LDS_SOFTMAX_BYTES + LDS_VMAX_BYTES + LDS_SCALE_TOTAL = LDS_SCALE_BYTES if const_expr(per_token_kv) else 0 + # Unique global symbol per compile to avoid clashes in a shared GPU context. + _smem_sym_name = ( + f"pa_ps_smallblk_smem_bs{block_size}_ql{query_length}" + f"_qgs{query_group_size}_tv{int(trans_v)}_qd{query_input_dtype}" + f"_ptkv{int(per_token_kv)}" + ) + allocator = SmemAllocator(None, arch=arch, global_sym_name=_smem_sym_name) + logits_off = 0 + allocator.ptr = LDS_LOGITS_BYTES + softmax_off = LDS_LOGITS_BYTES + allocator.ptr += LDS_SOFTMAX_TOTAL + # K/V per-token scale staging LDS (per_token_kv only). + scale_off_ps = softmax_off + LDS_SOFTMAX_TOTAL + allocator.ptr += LDS_SCALE_TOTAL + bt_off = scale_off_ps + LDS_SCALE_TOTAL + allocator.ptr += NUM_WARPS * TLOOP * 4 + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def pa_decode_ps_kernel( + exp_sums_ptr: fx.Tensor, + max_logits_ptr: fx.Tensor, + tmp_out_ptr: fx.Tensor, + query_ptr: fx.Tensor, + key_cache_ptr: fx.Tensor, + value_cache_ptr: fx.Tensor, + block_tables_ptr: fx.Tensor, + context_lengths_ptr: fx.Tensor, + key_scale_ptr: fx.Tensor, + value_scale_ptr: fx.Tensor, + stride_q_seq: Int32, + stride_q_head: Int32, + stride_k_block: Int32, + stride_k_head: Int32, + stride_v_block: Int32, + stride_v_head: Int32, + stride_es_seq: Int32, + stride_es_head: Int32, + stride_es_part: Int32, + stride_to_seq: Int32, + stride_to_head: Int32, + stride_to_part: Int32, + stride_to_group: Int32, + stride_bt_seq: Int32, + # Per-token K/V scale strides (per_token_kv only), scale layout + # [num_blocks, num_kv_heads, block_size]; both 0 for per-tensor. + stride_ks_block: Int32, + stride_ks_head: Int32, + ): + tid = fx.Int32(gpu.thread_id("x")) + batch_idx = fx.Int32(gpu.block_id("x")) + kv_h = fx.Int32(gpu.block_id("y")) + partition_idx = fx.Int32(gpu.block_id("z")) + + cl_global_ptr = _extract_global_ptr(context_lengths_ptr) + context_len = _global_load_i32(cl_global_ptr, batch_idx) + + lane16id = tid & fx.Int32(15) + rowid = (tid >> fx.Int32(4)) & fx.Int32(3) + warp_id = tid >> fx.Int32(6) + + q_rsrc = buffer_ops.create_buffer_resource(query_ptr, max_size=True) + k_global_ptr = _extract_global_ptr(key_cache_ptr) + v_global_ptr = _extract_global_ptr(value_cache_ptr) + bt_rsrc = buffer_ops.create_buffer_resource(block_tables_ptr, max_size=False) + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True) + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True) + to_rsrc = buffer_ops.create_buffer_resource(tmp_out_ptr, max_size=True) + ks_rsrc = buffer_ops.create_buffer_resource(key_scale_ptr, max_size=True) + vs_rsrc = buffer_ops.create_buffer_resource(value_scale_ptr, max_size=True) + + q_scale_val = arith.constant(1.0, type=T.f32) + # Per-tensor K/V scales load from index 0; per_token_kv stages per-token + # scales to LDS (see _stage_small_block_kv_scales). + if const_expr(per_token_kv): + k_scale_val = arith.constant(1.0, type=T.f32) + v_scale_val = arith.constant(1.0, type=T.f32) + else: + k_scale_val = buffer_ops.buffer_load( + ks_rsrc, arith.constant(0, type=T.i32), vec_width=1 + ) + v_scale_val = buffer_ops.buffer_load( + vs_rsrc, arith.constant(0, type=T.i32), vec_width=1 + ) + + smem_base = allocator.get_base() + logits_lds_i32 = SmemPtr( + smem_base, logits_off, T.i32, shape=(LDS_LOGITS_BYTES // 4,) + ).get() + softmax_lds_f32 = SmemPtr( + smem_base, softmax_off, T.f32, shape=(LDS_SOFTMAX_TOTAL // 4,) + ).get() + logits_lds_i64 = SmemPtr( + smem_base, logits_off, T.i64, shape=(LDS_LOGITS_BYTES // 8,) + ).get() + bt_lds_i32 = SmemPtr(smem_base, bt_off, T.i32, shape=(NUM_WARPS * TLOOP,)).get() + if const_expr(per_token_kv): + scale_lds_f32 = SmemPtr( + smem_base, scale_off_ps, T.f32, shape=(LDS_SCALE_BYTES // 4,) + ).get() + else: + scale_lds_f32 = None + + _softmax_scale_const = arith.constant(_softmax_scale, type=T.f32) + _softmax_q_scale = _softmax_scale_const * q_scale_val + _scale = _softmax_q_scale * k_scale_val + c_w = arith.constant(WARP_SIZE, type=T.i32) + NEG_INF = arith.constant(float("-inf"), type=T.f32) + ZERO_F = arith.constant(0.0, type=T.f32) + c_cps = arith.constant(KV_COMPUTE_BLOCK, type=T.i32) + c_query_group_size = arith.constant(query_group_size, type=T.i32) + + local_qhead_idx = warp_id * arith.constant(4, type=T.i32) + rowid + + ( + _k_tok_thread_base_unused, + _c_tok_stride_dw_unused, + _k_he_off_dw_unused, + _v_tok_thread_off, + _vhead_elem_dw, + _kv_tok_thread_base, + _prob_wr_thread_base, + _pv_prob_read_base, + _sm_max_off, + _sm_sum_off, + _sm_rd_max_offs, + _sm_rd_sum_offs, + _sm_vmax_wr_off, + _sm_vmax_rd_offs, + ) = _build_pa_thread_invariants( + warp_id, + lane16id, + rowid, + trans_v=trans_v, + per_token_kv=per_token_kv, + qkhe_loop=_QKHELOOP, + vhe_loop=_VHELOOP, + ) + + ( + _load_kv_scale_scalars_unused, + _load_v_and_scales_unused, + _store_vmax_warp, + _qk_and_intra_softmax, + _cross_warp_softmax_and_prob_pack, + _pv_mfma, + ) = _make_pa_phase_helpers( + trans_v=trans_v, + per_token_q=True, + per_token_kv=per_token_kv, + needs_mask=True, + query_length=query_length, + kv_h=kv_h, + v_global_ptr=v_global_ptr, + ks_rsrc=ks_rsrc, + vs_rsrc=vs_rsrc, + logits_lds_i32=logits_lds_i32, + logits_lds_i64=logits_lds_i64, + softmax_lds_f32=softmax_lds_f32, + scale_lds_f32=scale_lds_f32, + stride_ks_block=arith.constant(0, type=T.i32), + stride_ks_head=arith.constant(0, type=T.i32), + softmax_scale_base=_softmax_scale_const, + softmax_q_scale=_softmax_q_scale, + k_scale_val=k_scale_val, + scale=_scale, + v_scale_val=v_scale_val, + warp_id=warp_id, + lane16id=lane16id, + rowid=rowid, + k_tok_thread_base=_k_tok_thread_base_unused, + v_tok_thread_off=_v_tok_thread_off, + vhead_elem_dw=_vhead_elem_dw, + kv_tok_thread_base=_kv_tok_thread_base, + prob_wr_thread_base=_prob_wr_thread_base, + pv_prob_read_base=_pv_prob_read_base, + sm_max_off=_sm_max_off, + sm_sum_off=_sm_sum_off, + sm_rd_max_offs=_sm_rd_max_offs, + sm_rd_sum_offs=_sm_rd_sum_offs, + sm_vmax_wr_off=_sm_vmax_wr_off, + sm_vmax_rd_offs=_sm_vmax_rd_offs, + c_w=c_w, + neg_inf=NEG_INF, + zero_f=ZERO_F, + cache_scale_vecs=per_token_kv, + head_size=_HEAD, + qkhe_loop=_QKHELOOP, + vhe_loop=_VHELOOP, + ) + + def _store_partition_results(eqgs_lane, running_sum, running_max, outs_norm): + for vhe in range_constexpr(_VHELOOP): + hs_base = ( + fx.Int32(vhe * NUM_WARPS * MFMA_N) + + warp_id * fx.Int32(MFMA_N) + + rowid * fx.Int32(4) + ) + to_off = ( + batch_idx * stride_to_seq + + kv_h * stride_to_head + + partition_idx * stride_to_part + + eqgs_lane * stride_to_group + + hs_base + ) + out_bf16 = fx.Vector(outs_norm[vhe]).to(fx.BFloat16) + buffer_ops.buffer_store(out_bf16, to_rsrc, to_off) + es_off = ( + batch_idx * stride_es_seq + + kv_h * stride_es_head + + partition_idx * stride_es_part + + eqgs_lane + ) + buffer_ops.buffer_store(fx.Float32(running_sum), es_rsrc, es_off) + buffer_ops.buffer_store(fx.Float32(running_max), ml_rsrc, es_off) + + # Slot covers >=1 contiguous 256-token sub-partitions; the inner scf.for + # walks them with online-softmax loop-carried state. + c_max_parts = arith.constant(max_context_partition_num, type=T.i32) + num_total_partitions = (context_len + c_cps - fx.Int32(1)) >> fx.Int32(8) + page_size_partitions = ( + num_total_partitions + c_max_parts - fx.Int32(1) + ) // c_max_parts + local_partition_start = partition_idx * page_size_partitions + local_partition_end_raw = (partition_idx + fx.Int32(1)) * page_size_partitions + local_partition_end = arith.select( + local_partition_end_raw < num_total_partitions, + local_partition_end_raw, + num_total_partitions, + ) + + def _unwrap(v): + return v.ir_value() if hasattr(v, "ir_value") else v + + # Loop state: _mtp_groups accumulators (rmax, rsum, outs...) + the current + # sub-partition's K/V tiles (loop-carried ping-pong for prefetch). + state_width = 2 + _VHELOOP + + def _pack_states(states, k_flat, v_flat): + flat = [] + for st in states: + rmax, rsum = st[0], st[1] + outs = [st[2 + vhe] for vhe in range_constexpr(_VHELOOP)] + flat.extend([_unwrap(rmax), _unwrap(rsum)]) + flat.extend(_unwrap(out) for out in outs) + flat.extend(_unwrap(v) for v in k_flat) + flat.extend(_unwrap(v) for v in v_flat) + return flat + + def _unpack_states(flat): + base = state_width * _mtp_groups + states = [ + tuple(flat[state_width * i + j] for j in range_constexpr(state_width)) + for i in range_constexpr(_mtp_groups) + ] + k_flat = list(flat[base : base + _N_K_h]) + v_flat = list(flat[base + _N_K_h : base + _N_K_h + _N_V_FLAT_h]) + return states, k_flat, v_flat + + init_states = [ + tuple( + [NEG_INF, ZERO_F] + + [ + arith.constant_vector(0.0, T.f32x4) + for _ in range_constexpr(_VHELOOP) + ] + ) + for _ in range(_mtp_groups) + ] + + loop_start = fx.Index(arith.unwrap(local_partition_start)) + loop_end = fx.Index(arith.unwrap(local_partition_end)) + loop_step = arith.index(1) + last_partition_idx = local_partition_end - fx.Int32(1) + + def _ptr8_to_v4i32(ptr8_val): + """`ptr addrspace(8)` descriptor → `<4 x i32>` via ptrtoint(i128) + + bitcast: a 128-bit type-pun, zero instructions (stays in SGPRs).""" + from flydsl._mlir import ir as _ir + from flydsl._mlir.dialects import llvm as _llvm + + i128_ty = _ir.IntegerType.get_signless(128) + v4i32_ty = _ir.VectorType.get([4], _ir.IntegerType.get_signless(32)) + i128_val = _llvm.ptrtoint(i128_ty, ptr8_val) + return _llvm.bitcast(v4i32_ty, i128_val) + + bt_rsrc_v4 = _ptr8_to_v4i32(bt_rsrc) + + def _s_buffer_load(soffset_bytes_i32, vec_width: int): + """Scalar buffer load (s_buffer_load_dword[x4]) -> SGPR. REQUIRES + soffset_bytes_i32 wave-uniform. Saves the vmcnt(0) drain + readfirstlane.""" + from flydsl._mlir import ir as _ir + from flydsl._mlir.dialects import llvm as _llvm + from flydsl.expr.rocdl import _to_ir as _rocdl_to_ir + + i32_ty = _ir.IntegerType.get_signless(32) + if const_expr(vec_width == 1): + result_type = i32_ty + suffix = "i32" + elif const_expr(vec_width == 4): + result_type = _ir.VectorType.get([4], i32_ty) + suffix = "v4i32" + else: + raise ValueError(f"_s_buffer_load: unsupported vec_width={vec_width}") + cache_policy = arith.constant(0, type=T.i32) + return _llvm.call_intrinsic( + result_type, + f"llvm.amdgcn.s.buffer.load.{suffix}", + [ + _rocdl_to_ir(bt_rsrc_v4), + _rocdl_to_ir(soffset_bytes_i32), + _rocdl_to_ir(cache_policy), + ], + [], + [], + ) + + def _pa_small_block_stage_phys_blocks(partition_block_base): + # bt offset is wave-uniform -> s_buffer_load into SGPRs, avoiding the + # vmcnt(0) drain and the downstream readfirstlane of the VMEM path. + if const_expr(block_size == 64): + bt_elem_off = batch_idx * stride_bt_seq + partition_block_base + warp_id + phys_blocks = _s_buffer_load(bt_elem_off * fx.Int32(4), vec_width=1) + else: + bt_elem_off = ( + batch_idx * stride_bt_seq + + partition_block_base + + warp_id * fx.Int32(TLOOP) + ) + phys_blocks = _s_buffer_load(bt_elem_off * fx.Int32(4), vec_width=TLOOP) + return phys_blocks + + def _pa_small_block_store_phys_blocks_to_lds(phys_block_vec): + if (lane16id | rowid) == fx.Int32(0): + if const_expr(block_size == 64): + # Each warp writes 1 i32 to bt_lds_i32[warp_id]; readers pull vec4 at 0. + fx.Vector.from_elements([phys_block_vec], dtype=fx.Int32).store( + bt_lds_i32, + [fx.Index(warp_id)], + ) + else: + phys_block_vec.store( + bt_lds_i32, + [fx.Index(warp_id * fx.Int32(TLOOP))], + ) + + def _pa_small_block_load_v_phys_blocks_from_lds(): + v_phys_blocks = [] + if const_expr(block_size == 64): + phys_block_vec = fx.Vector.load( + T.vec(VTLOOP, T.i32), bt_lds_i32, [fx.Index(0)] + ) + for vt in range_constexpr(VTLOOP): + v_phys_blocks.append(phys_block_vec[vt]) + else: + for vt in range_constexpr(VTLOOP): + bt_lds_off = fx.Int32(vt * TLOOP) + rowid + phys_block = fx.Vector.load( + T.vec(1, T.i32), bt_lds_i32, [fx.Index(bt_lds_off)] + )[0] + v_phys_blocks.append(phys_block) + return v_phys_blocks + + # Pre-load the FIRST (reverse-order = last) sub-partition's block-table + # entries before Q setup so the dependent K prefetch avoids table latency. + # Empty-slot guard: clamp to 0 so 0-iter CTAs' prologue reads stay in-bounds. + _safe_init_partition = arith.select( + local_partition_start < num_total_partitions, + last_partition_idx, + arith.constant(0, type=T.i32), + ) + first_block_base = _safe_init_partition * fx.Int32(_blocks_per_partition) + first_phys_blocks = _pa_small_block_stage_phys_blocks(first_block_base) + + # Pre-load Q for every MTP group ONCE before the KV loop (stays in registers). + q_frags_per_mtp = [] + qi_per_mtp = [] + qhi_per_mtp = [] + qscale_per_mtp = [] + for _mtp_g in range_constexpr(_mtp_groups): + mtp_prefetch = _prefetch_mtp_group_query( + q_rsrc, + batch_idx, + kv_h, + stride_q_seq, + stride_q_head, + lane16id, + local_qhead_idx, + mtp_group_idx=_mtp_g, + query_length=query_length, + query_group_size=query_group_size, + query_load_is_bf16=query_load_is_bf16, + q_lanes_per_head=_Q_LANES_PER_HEAD, + q_elems_per_lane=_Q_ELEMS_PER_LANE, + q_chunks_per_lane=_Q_CHUNKS_PER_LANE, + ) + qi_val, qhi_pos, q_frags, query_scale_lane = _finish_mtp_group_q_fragments( + logits_lds_i32, + logits_lds_i64, + softmax_lds_f32, + mtp_prefetch, + lane16id, + rowid, + local_qhead_idx, + head_size=_HEAD, + qkhe_loop=_QKHELOOP, + q_lanes_per_head=_Q_LANES_PER_HEAD, + q_elems_per_lane=_Q_ELEMS_PER_LANE, + ) + q_frags_per_mtp.append(q_frags) + qi_per_mtp.append(qi_val) + qhi_per_mtp.append(qhi_pos) + qscale_per_mtp.append(query_scale_lane) + + _pa_small_block_store_phys_blocks_to_lds(first_phys_blocks) + + # Per-token K/V scale staging (per_token_kv only): each thread stages its + # LDS slot t from that token's page. Scale layout [num_blocks, kv_heads, block_size]. + def _stage_small_block_kv_scales(): + t = warp_id * fx.Int32(WARP_SIZE) + rowid * fx.Int32(MFMA_N) + lane16id + part_page = _udiv_const(t, _block_size) + tok_in_page = _urem_const(t, _block_size) + phys = fx.Vector.load(T.vec(1, T.i32), bt_lds_i32, [fx.Index(part_page)])[0] + scale_idx = phys * stride_ks_block + kv_h * stride_ks_head + tok_in_page + k_scale_scalar = buffer_ops.buffer_load( + ks_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + v_scale_scalar = buffer_ops.buffer_load( + vs_rsrc, scale_idx, vec_width=1, dtype=fx.Float32 + ) + fx.Vector.from_elements([k_scale_scalar], dtype=fx.Float32).store( + scale_lds_f32, [fx.Index(t)] + ) + fx.Vector.from_elements([v_scale_scalar], dtype=fx.Float32).store( + scale_lds_f32, [fx.Index(fx.Int32(LDS_SCALE_V_OFFSET) + t)] + ) + + def _load_small_block_scale_vecs(): + k_scale_vecs = [] + v_scale_vecs = [] + for td in range_constexpr(TLOOP): + row = _kv_tok_thread_base + arith.constant(td * MFMA_N, type=T.i32) + k_scale_vecs.append( + vector.load_op(T.f32x4, scale_lds_f32, [fx.Index(row)]) + ) + v_scale_vecs.append( + vector.load_op( + T.f32x4, + scale_lds_f32, + [fx.Index(fx.Int32(LDS_SCALE_V_OFFSET) + row)], + ) + ) + return k_scale_vecs, v_scale_vecs + + # Pre-load the FIRST sub-partition's K so the body can prefetch the next + # K in parallel with the current QK MFMA. + k_flat0 = _pa_small_block_load_k_flat( + k_global_ptr, + kv_h, + stride_k_block, + stride_k_head, + lane16id, + rowid, + block_size=_block_size, + phys_blocks=first_phys_blocks, + qkhe_loop=_QKHELOOP, + ) + gpu.barrier() + # Prologue V load (ping-pong with K): issue iter 0's V here so the body can + # prefetch iter N+1's V behind the next QK MFMA. (barrier ensures LDS vis). + _v_phys_blocks0 = _pa_small_block_load_v_phys_blocks_from_lds() + _v_results0 = _pa_small_block_load_v_trans( + v_global_ptr, + kv_h, + stride_v_block, + stride_v_head, + warp_id, + lane16id, + rowid, + _v_phys_blocks0, + block_size=_block_size, + head_size=_HEAD, + vhe_loop=_VHELOOP, + ) + v_flat0 = _flatten_v_results(_v_results0, vhe_loop=_VHELOOP) + # GOTCHA: do NOT wrap this loop in `if _is_valid:`. ReplaceIfWithDispatch + # copies the if-body into a synthetic fn; the scf.for `ast.Yield` makes + # it a generator (never executes), leaving the scf.if then-region empty. + # Run unconditionally; empty slots iterate 0 times and yield init state. + for sub_part_ib, state in range( + loop_start, + loop_end, + loop_step, + init=_pack_states(init_states, k_flat0, v_flat0), + ): + cur_states, k_flat, v_flat = _unpack_states(state) + # Reverse iteration: walk from last_partition_idx down to + # local_partition_start (sink-prone partition 0 processed last). + _sub_raw_i32 = arith.index_cast(T.i32, sub_part_ib) + sub_part_i32 = last_partition_idx - (_sub_raw_i32 - local_partition_start) + sub_token_start = sub_part_i32 * c_cps + + # K and V come from loop-carried state (prefetched at prev iter end); + # their VMEM latency overlaps the prev PV MFMA / next QK+softmax. + k_ops = _unflatten_k(k_flat, qkhe_loop=_QKHELOOP) + v_results = _unflatten_v_results(v_flat, vhe_loop=_VHELOOP) + + # Per-token K/V scale staging (per_token_kv only): stage to LDS once + # per partition, read cached f32x4 vecs reused across all MTP groups. + if const_expr(per_token_kv): + _stage_small_block_kv_scales() + gpu.barrier() + k_scale_vecs, v_scale_vecs = _load_small_block_scale_vecs() + + # NEXT sub-partition's K base (reverse: sub_part_i32 - 1), clamped to + # local_partition_start so the final iter's prefetch stays in-window. + next_part_i32 = sub_part_i32 - fx.Int32(1) + next_safe_part = arith.select( + next_part_i32 >= local_partition_start, + next_part_i32, + local_partition_start, + ) + next_block_base = next_safe_part * fx.Int32(_blocks_per_partition) + + new_states = [] + k_next_flat = None + for _mtp_g in range_constexpr(_mtp_groups): + state = cur_states[_mtp_g] + rmax, rsum = state[0], state[1] + outs = [state[2 + vhe] for vhe in range_constexpr(_VHELOOP)] + causal_bound = ( + context_len + + arith.constant(1 - query_length, type=T.i32) + + qi_per_mtp[_mtp_g] + ) + + if const_expr(per_token_kv): + d_out, v_scales = _qk_and_intra_softmax( + k_ops, + sub_token_start, + q_frags_per_mtp[_mtp_g], + causal_bound, + query_scale_lane=qscale_per_mtp[_mtp_g], + preloaded_scales=(k_scale_vecs, v_scale_vecs), + ) + else: + d_out = _qk_and_intra_softmax( + k_ops, + sub_token_start, + q_frags_per_mtp[_mtp_g], + causal_bound, + query_scale_lane=qscale_per_mtp[_mtp_g], + ) + v_scales = None + + if const_expr(_mtp_g == _mtp_groups - 1): + next_phys_blocks = _pa_small_block_stage_phys_blocks( + next_block_base + ) + + # per_token_kv: stage cross-warp v_scale_max to LDS for + # _cross_warp_softmax_and_prob_pack's norm_factor. + if const_expr(per_token_kv): + _store_vmax_warp( + sub_token_start, seq_end=context_len, v_scale_vecs=v_scales + ) + + gpu.barrier() + + rmax, rsum, outs, v_correction = _cross_warp_softmax_and_prob_pack( + d_out, rmax, rsum, outs, v_scales + ) + + # Next K prefetch on the LAST MTP iter, after cross_warp softmax + # but BEFORE _pv_mfma, so K VMEM latency overlaps the PV MFMA. + if const_expr(_mtp_g == _mtp_groups - 1): + _pa_small_block_store_phys_blocks_to_lds(next_phys_blocks) + k_next_flat = _pa_small_block_load_k_flat( + k_global_ptr, + kv_h, + stride_k_block, + stride_k_head, + lane16id, + rowid, + block_size=_block_size, + phys_blocks=next_phys_blocks, + qkhe_loop=_QKHELOOP, + ) + gpu.barrier() + outs = _pv_mfma(v_results, outs, v_correction) + new_states.append(tuple([rmax, rsum] + outs)) + + # Cross-iter V prefetch (ping-pong): issue NEXT iter's V AFTER PV MFMA + # (current V vgprs now free); latency hidden behind next QK MFMA + softmax. + _v_phys_blocks_next = _pa_small_block_load_v_phys_blocks_from_lds() + _v_next_results = _pa_small_block_load_v_trans( + v_global_ptr, + kv_h, + stride_v_block, + stride_v_head, + warp_id, + lane16id, + rowid, + _v_phys_blocks_next, + block_size=_block_size, + head_size=_HEAD, + vhe_loop=_VHELOOP, + ) + v_next_flat = _flatten_v_results(_v_next_results, vhe_loop=_VHELOOP) + + results = yield _pack_states(new_states, k_next_flat, v_next_flat) + + # Normalize and store one output slot per MTP group. + final_states, _final_k_flat, _final_v_flat = _unpack_states(results) + for _mtp_g in range_constexpr(_mtp_groups): + final_state = final_states[_mtp_g] + rmax_raw, rsum_raw = final_state[0], final_state[1] + outs_raw = [final_state[2 + vhe] for vhe in range_constexpr(_VHELOOP)] + running_max = fx.Float32(rmax_raw) + running_sum = fx.Float32(rsum_raw) + outs = [fx.Vector(out_raw) for out_raw in outs_raw] + outs_norm = _normalize_pa_output(running_sum, outs, ZERO_F) + eqgs_lane = qi_per_mtp[_mtp_g] * c_query_group_size + qhi_per_mtp[_mtp_g] + _store_partition_results(eqgs_lane, running_sum, running_max, outs_norm) + + @flyc.jit + def launch_pa_decode_ps_small_block( + exp_sums: fx.Tensor, + max_logits: fx.Tensor, + tmp_out: fx.Tensor, + query: fx.Tensor, + key_cache: fx.Tensor, + value_cache: fx.Tensor, + block_tables: fx.Tensor, + context_lengths: fx.Tensor, + key_scale: fx.Tensor, + value_scale: fx.Tensor, + s_q_seq: Int32, + s_q_head: Int32, + s_k_block: Int32, + s_k_head: Int32, + s_v_block: Int32, + s_v_head: Int32, + s_es_seq: Int32, + s_es_head: Int32, + s_es_part: Int32, + s_to_seq: Int32, + s_to_head: Int32, + s_to_part: Int32, + s_to_group: Int32, + s_bt_seq: Int32, + s_ks_block: Int32, + s_ks_head: Int32, + gx: Int32, + gy: Int32, + gz: Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + pa_decode_ps_kernel( + exp_sums, + max_logits, + tmp_out, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + s_q_seq, + s_q_head, + s_k_block, + s_k_head, + s_v_block, + s_v_head, + s_es_seq, + s_es_head, + s_es_part, + s_to_seq, + s_to_head, + s_to_part, + s_to_group, + s_bt_seq, + s_ks_block, + s_ks_head, + ).launch(grid=(gx, gy, gz), block=(BLOCK_THREADS, 1, 1), stream=stream) + + return { + "launch": launch_pa_decode_ps_small_block, + "kernel": pa_decode_ps_kernel, + "allocator": allocator, + "mtp_groups": _mtp_groups, + } + + +@functools.lru_cache(maxsize=64) +def compile_pa_decode_ps_reduce( + *, + head_dim: int, + eqgs: int, + max_parts: int, + output_dtype_str: str = "bf16", + arch: str = "", +): + """Combine per-partition NORMALIZED partials into the final output. + + Compute kernel writes, per partition p (g in [0, eqgs)): + temporary_output[..,p,g,:] = (sum_t P[t]*V[t]) / sum_t P[t] (bf16, NORMALIZED) + exp_sums[..,p,g] = sum_t P[t] (f32) + max_logits[..,p,g] = max logit (f32) + This vLLM/Gluon NORMALIZED-partial convention differs from + `pa_decode_reduce`'s un-normalized-numerator contract, hence this dedicated + reduce. Merge: + gmax = max_p max_logits[p] + w[p] = exp2((max_logits[p] - gmax) * LOG2E) (logits are natural-domain) + out = (sum_p w[p] * exp_sums[p] * norm_out[p]) / sum_p w[p]*exp_sums[p] + Grid (batch, num_kv_heads); each thread walks <=max_parts partitions serially. + """ + if not arch: + arch = get_hip_arch() + _OUT_FX = {"bf16": fx.BFloat16, "f16": fx.Float16, "f32": fx.Float32}[ + output_dtype_str + ] + _HD = head_dim + _EQGS = eqgs + _MP = max_parts + # Each thread owns a contiguous _VEC-wide head-dim slice for ONE group g + # (stats/weights computed once per thread); _VEC divides _HD -> coalesced load. + _VEC = 4 if (head_dim % 4 == 0) else (2 if (head_dim % 2 == 0) else 1) + _DV = _HD // _VEC # vector-slots per group along head-dim + _N = _EQGS * _DV # total (g, d-slot) work items per (batch, kv_head) + + @flyc.kernel(known_block_size=(BLOCK_THREADS, 1, 1)) + def _reduce_kernel( + output_ptr: fx.Tensor, + exp_sums_ptr: fx.Tensor, + max_logits_ptr: fx.Tensor, + tmp_out_ptr: fx.Tensor, + stride_o_seq: Int32, + stride_o_head: Int32, + stride_es_seq: Int32, + stride_es_head: Int32, + stride_es_part: Int32, + stride_to_seq: Int32, + stride_to_head: Int32, + stride_to_part: Int32, + stride_to_group: Int32, + num_kv_heads: Int32, + ) -> None: + tid = fx.Int32(gpu.thread_id("x")) + batch_idx = fx.Int32(gpu.block_id("x")) + kv_h = fx.Int32(gpu.block_id("y")) + + o_rsrc = buffer_ops.create_buffer_resource(output_ptr, max_size=True) + es_rsrc = buffer_ops.create_buffer_resource(exp_sums_ptr, max_size=True) + ml_rsrc = buffer_ops.create_buffer_resource(max_logits_ptr, max_size=True) + to_rsrc = buffer_ops.create_buffer_resource(tmp_out_ptr, max_size=True) + + c_zero = arith.constant(0.0, type=T.f32) + c_one = arith.constant(1.0, type=T.f32) + c_neginf = arith.constant(float("-inf"), type=T.f32) + + es_base = batch_idx * stride_es_seq + kv_h * stride_es_head + to_base = batch_idx * stride_to_seq + kv_h * stride_to_head + o_base = batch_idx * stride_o_seq + kv_h * stride_o_head + + _out_vec_ty = T.vec(_VEC, _OUT_FX.ir_type) + # Thread owns work items n = tid, tid+BLOCK_THREADS, ... over (g, d-slot): + # n -> (g, dv), dv the _VEC-wide slot, d0 = dv*_VEC. + for _i in range_constexpr((_N + BLOCK_THREADS - 1) // BLOCK_THREADS): + n = tid + fx.Int32(_i * BLOCK_THREADS) + if const_expr((_N % BLOCK_THREADS) != 0): + do = n < fx.Int32(_N) + g = n // fx.Int32(_DV) + dv = n % fx.Int32(_DV) + d0 = dv * fx.Int32(_VEC) + + # Per-partition stats depend only on g, loaded once per thread. + # Pass 1: global max. Pass 2: wsum[p] and running denom gsum. + gmax = fx.Float32(c_neginf) + mls = [] + for p in range_constexpr(_MP): + ml = buffer_ops.buffer_load( + ml_rsrc, + es_base + fx.Int32(p) * stride_es_part + g, + vec_width=1, + dtype=T.f32, + ) + mls.append(ml) + gmax = _maxnumf(gmax, fx.Float32(ml)) + gmax_ok = arith.select( + arith.unwrap(gmax) > c_neginf, arith.unwrap(gmax), c_zero + ) + + wsums = [] + gsum = fx.Float32(c_zero) + for p in range_constexpr(_MP): + es = buffer_ops.buffer_load( + es_rsrc, + es_base + fx.Int32(p) * stride_es_part + g, + vec_width=1, + dtype=T.f32, + ) + # max_logits[p] is natural-domain; the merge MUST use the same + # exp2(diff * LOG2E) factor as the compute kernel's softmax. + w = _exp2_f32_fast( + fx.Float32( + arith.mulf( + arith.subf(arith.unwrap(fx.Float32(mls[p])), gmax_ok), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + w = arith.select( + arith.unwrap(fx.Float32(mls[p])) > c_neginf, arith.unwrap(w), c_zero + ) + wsum = arith.mulf(w, arith.unwrap(fx.Float32(es))) + wsums.append(wsum) + gsum = fx.Float32(arith.addf(arith.unwrap(gsum), wsum)) + + safe = arith.select(arith.unwrap(gsum) > c_zero, arith.unwrap(gsum), c_one) + inv = arith.unwrap(_rcp_f32(fx.Float32(safe))) + + # WARNING: temporary_output is ALWAYS bf16 (compute kernel writes bf16 + # partials regardless of `output` dtype), so the load dtype MUST be bf16. + accs = [c_zero] * _VEC + for p in range_constexpr(_MP): + to_off = ( + to_base + fx.Int32(p) * stride_to_part + g * stride_to_group + d0 + ) + to_vec = buffer_ops.buffer_load( + to_rsrc, to_off, vec_width=_VEC, dtype=fx.BFloat16 + ) + for c in range_constexpr(_VEC): + tv = vector.extract( + to_vec, static_position=[c], dynamic_position=[] + ) + tv_f = arith.extf(T.f32, tv) + accs[c] = arith.addf(accs[c], arith.mulf(wsums[p], tv_f)) + + out_vec = arith.constant_vector(0.0, _out_vec_ty) + for c in range_constexpr(_VEC): + ov = _OUT_FX(arith.mulf(accs[c], inv)) + out_vec = vector.insert( + arith.unwrap(ov), out_vec, static_position=[c], dynamic_position=[] + ) + o_off = o_base + g * fx.Int32(_HD) + d0 + if const_expr((_N % BLOCK_THREADS) != 0): + if do: + buffer_ops.buffer_store(out_vec, o_rsrc, o_off) + else: + buffer_ops.buffer_store(out_vec, o_rsrc, o_off) + + @flyc.jit + def _launcher( + output_ptr, + exp_sums_ptr, + max_logits_ptr, + tmp_out_ptr, + stride_o_seq, + stride_o_head, + stride_es_seq, + stride_es_head, + stride_es_part, + stride_to_seq, + stride_to_head, + stride_to_part, + stride_to_group, + num_kv_heads, + grid_b, + grid_h, + stream: fx.Stream = fx.Stream(None), + ): + _reduce_kernel( + output_ptr, + exp_sums_ptr, + max_logits_ptr, + tmp_out_ptr, + stride_o_seq, + stride_o_head, + stride_es_seq, + stride_es_head, + stride_es_part, + stride_to_seq, + stride_to_head, + stride_to_part, + stride_to_group, + num_kv_heads, + ).launch(grid=(grid_b, grid_h, 1), block=(BLOCK_THREADS, 1, 1), stream=stream) + + return {"launch": _launcher, "kernel": _reduce_kernel} + + +def pa_decode_ps_launch( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + context_lengths: torch.Tensor, + softmax_scale: float, + key_scale: torch.Tensor = None, + value_scale: torch.Tensor = None, + *, + block_tables: torch.Tensor = None, # [num_seqs, max_blocks_per_seq] i32 + max_context_partition_num: int = 0, + exp_sums: torch.Tensor = None, + max_logits: torch.Tensor = None, + temporary_output: torch.Tensor = None, + stream=None, +) -> str: + """Launch the small-block (block_size 16/64) persistent-scheduling PA decode. + MSLK port: only the small-block compute path (see module docstring). + """ + num_query_heads = query.shape[1] + num_kv_heads = key_cache.shape[1] + trans_v = len(value_cache.shape) == 5 + query_input_dtype = _get_query_input_dtype(query) + + dev = query.device + is_graph_capturing = _is_current_stream_capturing() + + key_scale = _prepare_scale_tensor( + "key_scale", + key_scale, + device=dev, + is_graph_capturing=is_graph_capturing, + ) + value_scale = _prepare_scale_tensor( + "value_scale", + value_scale, + device=dev, + is_graph_capturing=is_graph_capturing, + ) + if query_input_dtype == "packed_fp8": + raise ValueError( + "`pa_decode_ps_launch` no longer accepts host query_scale and only supports " + "bf16/f16 query inputs with kernel-internal query scale computation." + ) + + # per-token K/V path iff scale tensor is >1-D (one scale per (block,head,token)). + per_token_kv = key_scale.ndim > 1 + + query_length = query.shape[0] // context_lengths.shape[0] + query_group_size = num_query_heads // num_kv_heads + + # Strides for key_scale/value_scale + if per_token_kv: + stride_ks_block = key_scale.stride(0) + stride_ks_head = key_scale.stride(1) + else: + stride_ks_block = 0 + stride_ks_head = 0 + + s = stream or torch.cuda.current_stream() + + # Key cache shape: [num_blocks, num_kv_heads, head_size // 16, block_size, 16]. + block_size = key_cache.shape[-2] + if block_size not in _PA_DECODE_PS_SMALL_BLOCK_SIZES: + raise NotImplementedError( + f"pa_decode_ps_launch (MSLK port): only small block_size " + f"{_PA_DECODE_PS_SMALL_BLOCK_SIZES} is supported; got block_size={block_size}. " + "The metadata/aiter and sliding-window paths are not ported." + ) + if block_tables is None: + raise ValueError( + f"pa_decode_ps_launch: block_size={block_size} requires `block_tables` " + "(per-sequence physical block index table)." + ) + batch_size = context_lengths.shape[0] + head_size = query.shape[-1] + eqgs = query_length * query_group_size + context_partition_size = KV_COMPUTE_BLOCK + blocks_per_partition = context_partition_size // block_size + if max_context_partition_num == 0: + max_context_partition_num = get_recommended_splits( + batch_size, + num_kv_heads, + split_kv_blocks=blocks_per_partition, + ) + if is_graph_capturing and ( + exp_sums is None or max_logits is None or temporary_output is None + ): + raise ValueError( + "CUDA graph capture requires preallocated `exp_sums`, `max_logits`, " + "and `temporary_output` for the small-block PS path." + ) + if exp_sums is None: + exp_sums = torch.zeros( + batch_size, + num_kv_heads, + max_context_partition_num, + eqgs, + device=dev, + dtype=torch.float32, + ) + if max_logits is None: + max_logits = torch.full( + (batch_size, num_kv_heads, max_context_partition_num, eqgs), + float("-inf"), + device=dev, + dtype=torch.float32, + ) + if temporary_output is None: + temporary_output = torch.zeros( + batch_size, + num_kv_heads, + max_context_partition_num, + eqgs, + head_size, + device=dev, + dtype=torch.bfloat16, + ) + compiled_small = compile_pa_decode_ps( + block_size=block_size, + max_context_partition_num=max_context_partition_num, + softmax_scale=softmax_scale, + trans_v=trans_v, + query_group_size=query_group_size, + per_token_kv=per_token_kv, + query_length=query_length, + query_input_dtype=query_input_dtype, + head_dim=int(head_size), + ) + output_5d = output.reshape( + batch_size, query_length, num_kv_heads, query_group_size, head_size + ) + compiled_small["launch"]( + exp_sums, + max_logits, + temporary_output, + query, + key_cache, + value_cache, + block_tables, + context_lengths, + key_scale, + value_scale, + query.stride(0), + query.stride(1), + key_cache.stride(0), + key_cache.stride(1), + value_cache.stride(0), + value_cache.stride(1), + exp_sums.stride(0), + exp_sums.stride(1), + exp_sums.stride(2), + temporary_output.stride(0), + temporary_output.stride(1), + temporary_output.stride(2), + temporary_output.stride(3), + block_tables.stride(0), + stride_ks_block, + stride_ks_head, + batch_size, + num_kv_heads, + max_context_partition_num, + s, + ) + + # Dedicated reduce for NORMALIZED partials (pa_decode_reduce expects + # un-normalized numerators, so it can't be reused). + out_dtype_str = _get_output_dtype_str(output) + # Reduce needs a contiguous (eqgs, head_size) block per (batch, kv_head), + # which holds only for query_length == 1. + if query_length != 1: + raise NotImplementedError( + "pa_decode_ps_launch reduce: query_length > 1 not supported yet " + f"(got query_length={query_length})." + ) + reduce_compiled = compile_pa_decode_ps_reduce( + head_dim=int(head_size), + eqgs=int(eqgs), + max_parts=int(max_context_partition_num), + output_dtype_str=out_dtype_str, + ) + # output_5d: [batch, 1, num_kv_heads, query_group_size, head_size]. + reduce_compiled["launch"]( + output_5d, + exp_sums, + max_logits, + temporary_output, + num_kv_heads * eqgs * head_size, # stride_o_seq (one batch element) + eqgs * head_size, # stride_o_head (one kv head within batch) + exp_sums.stride(0), + exp_sums.stride(1), + exp_sums.stride(2), + temporary_output.stride(0), + temporary_output.stride(1), + temporary_output.stride(2), + temporary_output.stride(3), + num_kv_heads, + batch_size, + num_kv_heads, + s, # same stream as the compute launch — required for CUDA-graph capture + ) + return "ps_small_block" + + +# ── AOT interface ───────────────────────────────────────────────────────────── +# +# gfx950-only. softmax_scale is baked into the kernel, so AOT precompiles with the +# default 1/sqrt(head_dim); non-default scales JIT-compile on first use. + +AOT_ARCHS: List[str] = ["gfx950"] + +# Baked params match the cache builder (dense_kv_to_fp8_paged): block_size=16, +# per_token_kv, trans_v; qgs covers MQA(1)+GQA; max_parts is get_recommended_splits' {4, 8}. +_FP8_HEAD_SIZES = (128, 256) +_FP8_QGS = (1, 2, 4, 8, 16) +_FP8_MAX_PARTS = (4, 8) +_FP8_Q_DTYPES = ("bf16", "f16") + +AOT_CONFIGS: List[Dict[str, Any]] = [ + { + "head_dim": hd, + "query_group_size": qgs, + "max_context_partition_num": mp, + "query_input_dtype": qdt, + } + for hd in _FP8_HEAD_SIZES + for qgs in _FP8_QGS + for mp in _FP8_MAX_PARTS + for qdt in _FP8_Q_DTYPES +] + + +def compile_aot_config(config: Dict[str, Any], arch: str) -> None: + """Precompile one fp8 config (compute kernel + reduce). gfx950 only.""" + if not arch.startswith("gfx950"): + return + hd = config["head_dim"] + qgs = config["query_group_size"] + mp = config["max_context_partition_num"] + qdt = config["query_input_dtype"] + + compile_pa_decode_ps( + block_size=16, + max_context_partition_num=mp, + softmax_scale=1.0 / (hd**0.5), + trans_v=True, + query_group_size=qgs, + per_token_kv=True, + query_length=1, + query_input_dtype=qdt, + head_dim=hd, + ) + # Reduce output dtype = the query dtype (decode writes the query's dtype). + compile_pa_decode_ps_reduce( + head_dim=hd, + eqgs=qgs, # query_length (1) * query_group_size + max_parts=mp, + output_dtype_str=qdt, + arch=arch, + ) diff --git a/mslk/attention/fmha/flydsl/pa_decode_fp8_dispatch.py b/mslk/attention/fmha/flydsl/pa_decode_fp8_dispatch.py new file mode 100644 index 00000000..c7876f11 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_fp8_dispatch.py @@ -0,0 +1,154 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Public dispatcher for the FlyDSL native-fp8 paged-attention decode. + +Native-fp8 paged KV with a symmetric per-token scale (vLLM/Gluon layout); distinct +from the Triton int32-packed asymmetric-scale path, hence a separate guarded entry. + +Expected inputs (same CUDA device): + * query : [num_seqs, num_query_heads, head_size] bf16/f16. + * key_cache : [num_blocks, num_kv_heads, head_size // 16, block_size, 16] fp8. + * value_cache : [num_blocks, num_kv_heads, block_size // 16, head_size, 16] fp8 (transposed). + * key_scale/value_scale : per-token f32 [num_blocks, num_kv_heads, block_size, 1]. + * block_tables : [num_seqs, max_blocks_per_seq] int32. + * context_lengths : [num_seqs] int32. + +Only block_size in {16, 64}; query_length must be 1 (decode); gfx950 only. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +from mslk.flydsl.common import is_flydsl_available, require_flydsl + + +def is_fp8_paged_decode_available() -> bool: + """True when the FlyDSL native-fp8 paged decode can run on this arch (gfx950).""" + if not is_flydsl_available(): + return False + try: + from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] + + return get_rocm_arch().startswith("gfx950") + except Exception: + return False + + +def csr_to_block_tables( + kv_page_indices: torch.Tensor, # [total_pages] int32 — flat physical page ids + kv_indptr: torch.Tensor, # [num_seqs + 1] int32 — prefix sum of pages/seq +) -> torch.Tensor: + """Convert ragged CSR paging (kv_page_indices/kv_indptr) into a dense padded + block_tables[num_seqs, max_blocks_per_seq]. Rows right-padded with 0 (inert: + the walk is bounded by context_lengths). + """ + if kv_indptr.dtype != torch.int32: + kv_indptr = kv_indptr.to(torch.int32) + if kv_page_indices.dtype != torch.int32: + kv_page_indices = kv_page_indices.to(torch.int32) + dev = kv_page_indices.device + indptr = kv_indptr.to(torch.long) + num_seqs = indptr.numel() - 1 + counts = indptr[1:] - indptr[:-1] # pages per sequence + max_blocks = int(counts.max().item()) if num_seqs > 0 else 0 + max_blocks = max(max_blocks, 1) + block_tables = torch.zeros((num_seqs, max_blocks), dtype=torch.int32, device=dev) + for b in range(num_seqs): + lo = int(indptr[b].item()) + hi = int(indptr[b + 1].item()) + n = hi - lo + if n > 0: + block_tables[b, :n] = kv_page_indices[lo:hi] + return block_tables + + +def paged_attention_decode_fp8_csr( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + context_lengths: torch.Tensor, + kv_page_indices: torch.Tensor, # [total_pages] int32 + kv_indptr: torch.Tensor, # [num_seqs + 1] int32 + softmax_scale: float, + key_scale: torch.Tensor, + value_scale: torch.Tensor, + *, + max_context_partition_num: int = 0, + exp_sums: Optional[torch.Tensor] = None, + max_logits: Optional[torch.Tensor] = None, + temporary_output: Optional[torch.Tensor] = None, + stream: Optional[object] = None, +) -> str: + """CSR-paging entry: converts ragged kv_page_indices/kv_indptr then dispatches to + paged_attention_decode_fp8.""" + block_tables = csr_to_block_tables(kv_page_indices, kv_indptr) + return paged_attention_decode_fp8( + output, + query, + key_cache, + value_cache, + context_lengths, + block_tables, + softmax_scale, + key_scale, + value_scale, + max_context_partition_num=max_context_partition_num, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + stream=stream, + ) + + +def paged_attention_decode_fp8( + output: torch.Tensor, + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + context_lengths: torch.Tensor, + block_tables: torch.Tensor, + softmax_scale: float, + key_scale: torch.Tensor, + value_scale: torch.Tensor, + *, + max_context_partition_num: int = 0, + exp_sums: Optional[torch.Tensor] = None, + max_logits: Optional[torch.Tensor] = None, + temporary_output: Optional[torch.Tensor] = None, + stream: Optional[object] = None, +) -> str: + """Run the FlyDSL native-fp8 paged decode (writes into `output`). Guarded wrapper + around pa_decode_fp8.pa_decode_ps_launch; raises when FlyDSL/arch unavailable.""" + require_flydsl() + if not is_fp8_paged_decode_available(): + raise RuntimeError( + "FlyDSL native-fp8 paged decode requires gfx950 (CDNA4). " + "For the int32-packed Triton fp8 format, use triton_splitk.FwOp instead." + ) + from .pa_decode_fp8 import pa_decode_ps_launch + + return pa_decode_ps_launch( + output, + query, + key_cache, + value_cache, + context_lengths, + softmax_scale, + key_scale=key_scale, + value_scale=value_scale, + block_tables=block_tables, + max_context_partition_num=max_context_partition_num, + exp_sums=exp_sums, + max_logits=max_logits, + temporary_output=temporary_output, + stream=stream, + ) diff --git a/mslk/attention/fmha/flydsl/pa_decode_generic.py b/mslk/attention/fmha/flydsl/pa_decode_generic.py new file mode 100644 index 00000000..922b6164 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_generic.py @@ -0,0 +1,662 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL paged-attention decode (generic) — arch-generic fallback. + +Each warp owns one Q head (NUM_WARPS=4 heads/CTA), so softmax is intra-warp only. +Uses mfma_f32_16x16x32 (K=32); TLOOP covers TILE_N=64 tokens/step via 4 sub-tiles of 16. + +MFMA layout (mfma_f32_16x16x32_f16, wave64), lane l: + A/B: vec<8,f16>/lane → [row/col=l%16, k=(l//16)*8 : +8] + C: vec<4,f32>/lane → C[(l//16)*4+elem, l%16] +Per warp: tok_qk = lane%16 (token), k_grp = lane//16 (0..3, D chunk). +""" + +from __future__ import annotations + +import functools +from typing import Any, Optional + +import flydsl.compiler as flyc # pyre-ignore[21] +import flydsl.expr as fx # pyre-ignore[21] +import torch +from flydsl.expr import ( # pyre-ignore[21] + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr # pyre-ignore[21] + +from .pa_decode_reduce import pa_decode_reduce +from .utils import ( + dpp_xor_f32, + exp2_f32 as _exp2_fast, + maxnumf as _mxf, + rcp_f32, + smem_bytes, + WARP_SIZE, +) + +NUM_WARPS = 4 # warps per CTA = Q heads per CTA +MFMA_N = 16 # tokens per MFMA call +MFMA_K = 32 # K-dim of mfma_f32_16x16x32_f16 +TLOOP = NUM_WARPS # sub-tiles per step (each warp covers NUM_WARPS×16 = 64 tokens) +TILE_N = TLOOP * MFMA_N # 64 tokens per tile step +BLOCK = NUM_WARPS * WARP_SIZE # 256 threads + +_FX_DTYPE = {"f32": fx.Float32, "f16": fx.Float16, "bf16": fx.BFloat16} +LOG2E: float = 1.4426950408889634 + + +@functools.lru_cache(maxsize=256) +def compile_pa_decode_generic( + *, + head_size: int, + kv_dtype_str: str, + output_dtype_str: str, + split_k: int = 1, + arch: str = "", +) -> Any: # pyre-ignore[3] + if not arch: + arch = get_rocm_arch() + + assert head_size % MFMA_K == 0, f"head_size must be multiple of {MFMA_K}" + assert kv_dtype_str in ("f16", "bf16") + + _HEAD = head_size + _SK = split_k + _SPLIT = _SK > 1 + _FX_KV = _FX_DTYPE[kv_dtype_str] + _FX_OUT = _FX_DTYPE[output_dtype_str] + # MFMA intrinsic MUST match KV operand dtype (bf16 fails the _f16 verifier). + _mfma = ( + rocdl.mfma_f32_16x16x32_bf16 + if kv_dtype_str == "bf16" + else rocdl.mfma_f32_16x16x32_f16 + ) + _QK_GROUPS = _HEAD // MFMA_K # D/32 groups for Q·K + _PV_GROUPS = _HEAD // MFMA_N # D/16 groups for P·V + + # p_lds[NUM_WARPS*TILE_N]: P weights, [warp_id*TILE_N + td*MFMA_N + tok_qk]. + # No ms_lds/pv_lds — softmax and PV accum are intra-warp per Q head. + _P_ELEMS = NUM_WARPS * TILE_N + _LDS_TOTAL = _P_ELEMS * 4 # 1024 bytes + + cap = smem_bytes(arch) + if _LDS_TOTAL > cap: + raise ValueError(f"LDS {_LDS_TOTAL}B > {arch!r} cap {cap}B") + + alloc = SmemAllocator( + None, + arch=arch, + global_sym_name=f"pa_generic_h{_HEAD}_{kv_dtype_str}_nw{NUM_WARPS}_sk{_SK}", + ) + alloc.ptr = _LDS_TOTAL + + @flyc.kernel(known_block_size=(BLOCK, 1, 1)) + def pa_decode_generic_kernel( + out_ptr: fx.Tensor, + partial_max_ptr: fx.Tensor, + partial_sum_ptr: fx.Tensor, + q_ptr: fx.Tensor, + k_ptr: fx.Tensor, + v_ptr: fx.Tensor, + seq_ptr: fx.Tensor, + stride_qb: fx.Int32, + stride_qg: fx.Int32, + stride_qh: fx.Int32, + stride_kb: fx.Int32, + stride_km: fx.Int32, + stride_kg: fx.Int32, + stride_kh: fx.Int32, + num_hq: fx.Int32, + num_g: fx.Int32, + kv_max: fx.Int32, + num_hkv: fx.Int32, + softmax_scale: fx.Float32, + split_total: fx.Int32, + ) -> None: + tid = gpu.thread_idx.x + warp_id = tid >> fx.Int32(6) + lane = tid & fx.Int32(63) + + # tok_qk = lane%16 (N-col/token), k_grp = lane//16 (D chunk, 0..3) + tok_qk = lane & fx.Int32(MFMA_N - 1) + k_grp = lane >> fx.Int32(4) + + # Grid → (b, g, hq_block, split_idx) + flat = fx.Int32(gpu.block_idx.x) + if const_expr(_SPLIT): + split_idx = flat % split_total + rest = flat // split_total + else: + split_idx = fx.Int32(0) + rest = flat + + n_hq_blocks = (num_hq + fx.Int32(NUM_WARPS - 1)) // fx.Int32(NUM_WARPS) + hq_block = rest % n_hq_blocks + rest2 = rest // n_hq_blocks + g_idx = rest2 % num_g + b_idx = rest2 // num_g + + # Each warp owns ONE Q head + hq_abs = hq_block * fx.Int32(NUM_WARPS) + warp_id + hkv_abs = hq_abs * num_hkv // num_hq + + q_base = b_idx * stride_qb + g_idx * stride_qg + hq_abs * stride_qh + kv_base = b_idx * stride_kb + g_idx * stride_kg + hkv_abs * stride_kh + + c_zero = arith.constant(0.0, type=T.f32) + c_one = arith.constant(1.0, type=T.f32) + c_neginf = arith.constant(float("-inf"), type=T.f32) + zero_v4 = arith.constant_vector(0.0, T.vec(4, T.f32)) + zero_v8h = arith.constant_vector(0.0, T.vec(8, _FX_KV.ir_type)) + + seq_rsrc = buffer_ops.create_buffer_resource(seq_ptr, max_size=True) + q_rsrc = buffer_ops.create_buffer_resource(q_ptr, max_size=True) + k_rsrc = buffer_ops.create_buffer_resource(k_ptr, max_size=True) + v_rsrc = buffer_ops.create_buffer_resource(v_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True) + pm_rsrc = buffer_ops.create_buffer_resource(partial_max_ptr, max_size=True) + ps_rsrc = buffer_ops.create_buffer_resource(partial_sum_ptr, max_size=True) + + seq_len = buffer_ops.buffer_load(seq_rsrc, b_idx, vec_width=1, dtype=T.i32) + t_full = arith.select(seq_len > fx.Int32(0), seq_len, kv_max) + if const_expr(_SPLIT): + chunk = (t_full + split_total - fx.Int32(1)) // split_total + t_start = split_idx * chunk + t_end_raw = (split_idx + fx.Int32(1)) * chunk + t_end = arith.select(t_end_raw < t_full, t_end_raw, t_full) + else: + t_start = fx.Int32(0) + t_end = t_full + + smem = alloc.get_base() + p_lds = SmemPtr(smem, 0, T.f32, shape=(_P_ELEMS,)).get() + + # Pre-load Q A-frags: Q[hq_abs, g*MFMA_K + k_grp*8 : +8] as vec<8,f16> + q_frags = [] + for g in range_constexpr(_QK_GROUPS): + q_off = q_base + fx.Int32(g * MFMA_K) + k_grp * fx.Int32(8) + q_frags.append( + buffer_ops.buffer_load(q_rsrc, q_off, vec_width=8, dtype=_FX_KV) + ) + + # State: running max, running sum, then PV accum (_PV_GROUPS×4 C-elems/lane) + _init_neg = arith.constant(float("-inf"), type=T.f32) + _init_zer = arith.constant(0.0, type=T.f32) + _N_PV = _PV_GROUPS * 4 + _init_state = [_init_neg, _init_zer] + [_init_zer] * _N_PV + + _t_s = fx.Index(t_start) + _t_e = fx.Index(t_end) + + for _tile_i, state in range(_t_s, _t_e, arith.index(TILE_N), init=_init_state): + running_max = fx.Float32(state[0]) + running_sum = fx.Float32(state[1]) + pv_scalars = [state[2 + i] for i in range(_N_PV)] + + tile_start = fx.Int32(arith.index_cast(T.i32, _tile_i)) + + # ── QK: prefetch K for all TLOOP sub-tiles, then MFMA ───────── + # sub-tile td covers tokens tile_start + td*MFMA_N + tok_qk + k_frags_all = [] + for td in range_constexpr(TLOOP): + tok_td = tile_start + fx.Int32(td * MFMA_N) + tok_qk + k_frags_td = [] + for g in range_constexpr(_QK_GROUPS): + k_off = ( + kv_base + + tok_td * stride_km + + fx.Int32(g * MFMA_K) + + k_grp * fx.Int32(8) + ) + k_frags_td.append( + buffer_ops.buffer_load(k_rsrc, k_off, vec_width=8, dtype=_FX_KV) + ) + k_frags_all.append(k_frags_td) + + rocdl.sched_barrier(0) + + qk_vecs = [] + for td in range_constexpr(TLOOP): + qk_acc = zero_v4 + for g in range_constexpr(_QK_GROUPS): + qk_acc = _mfma( + T.vec(4, T.f32), + [q_frags[g], k_frags_all[td][g], qk_acc, 0, 0, 0], + ) + qk_vecs.append(qk_acc) + + # ── Softmax (intra-warp, no LDS) ────────────────────────────── + # QK scalar = C[elem=0, col=tok_qk] per sub-tile + qk_vals = [] + for td in range_constexpr(TLOOP): + tok_td = tile_start + fx.Int32(td * MFMA_N) + tok_qk + qk_raw = vector.extract( + qk_vecs[td], static_position=[0], dynamic_position=[] + ) + qk_sc = arith.mulf(qk_raw, arith.unwrap(softmax_scale)) + head_ok = hq_abs < num_hq + tok_ok = tok_td < t_end + in_range = arith.andi(arith.unwrap(head_ok), arith.unwrap(tok_ok)) + qk_vals.append(fx.Float32(arith.select(in_range, qk_sc, c_neginf))) + + # Intra-warp max across TLOOP*16 tokens (all in registers). + tile_max = qk_vals[0] + for td in range_constexpr(1, TLOOP): + tile_max = _mxf(tile_max, qk_vals[td]) + # DPP butterfly over tok_qk within each 16-lane k_grp segment. + for sh in (8, 4, 2, 1): + tile_max = _mxf(tile_max, dpp_xor_f32(tile_max, sh)) + # NOTE: no cross-k_grp reduce needed. The MFMA already accumulates all + # k_grps into C[0,tok_qk], so every k_grp holds the SAME full Q·K scalar. + + new_max = _mxf(running_max, tile_max) + rescale = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf(arith.unwrap(running_max), arith.unwrap(new_max)), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + + # P values, normalized by new_max (standard online softmax). + safe_max = fx.Float32( + arith.select( + arith.unwrap(new_max) > c_neginf, arith.unwrap(new_max), c_zero + ) + ) + p_vals = [] + intra_sum = fx.Float32(c_zero) + for td in range_constexpr(TLOOP): + tok_td = tile_start + fx.Int32(td * MFMA_N) + tok_qk + head_ok = hq_abs < num_hq + tok_ok = tok_td < t_end + in_range = arith.andi(arith.unwrap(head_ok), arith.unwrap(tok_ok)) + p_c = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf( + arith.unwrap(qk_vals[td]), arith.unwrap(safe_max) + ), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + p_c = fx.Float32(arith.select(in_range, arith.unwrap(p_c), c_zero)) + p_vals.append(p_c) + intra_sum = fx.Float32( + arith.addf(arith.unwrap(intra_sum), arith.unwrap(p_c)) + ) + + # DPP sum over tok_qk within 16-lane segment. + for sh in (8, 4, 2, 1): + intra_sum = fx.Float32( + arith.addf( + arith.unwrap(intra_sum), + arith.unwrap(dpp_xor_f32(intra_sum, sh)), + ) + ) + tile_sum = intra_sum + + new_sum = fx.Float32( + arith.addf( + arith.mulf(arith.unwrap(rescale), arith.unwrap(running_sum)), + arith.unwrap(tile_sum), + ) + ) + + # Write P to p_lds[warp_id*TILE_N + td*MFMA_N + tok_qk]. + for td in range_constexpr(TLOOP): + p_slot = fx.Index( + warp_id * fx.Int32(TILE_N) + fx.Int32(td * MFMA_N) + tok_qk + ) + vector.store( + fx.Vector.from_elements( + [arith.unwrap(p_vals[td])], dtype=fx.Float32 + ), + p_lds, + [p_slot], + ) + gpu.barrier() + + # ── PV MFMA ────────────────────────────────────────────────── + # 64 tokens, K=32 → 2 MFMA calls (halves 0..31, 32..63) per D-group. + # P B-frag: p_lds[warp*64 + k_grp*8 + j]. + # V A-frag: V[tok=tile_start + half*32 + k_grp*8+j, d_out=g*MFMA_N+tok_qk]. + + # Prefetch P frags for both halves + p_half_base = [ + warp_id * fx.Int32(TILE_N), + warp_id * fx.Int32(TILE_N) + fx.Int32(TILE_N // 2), + ] + p_frags = [] + for half in range_constexpr(2): + p_frag = zero_v8h + pbase = p_half_base[half] + k_grp * fx.Int32(8) + for j in range_constexpr(8): + pf_j = fx.Vector.load( + T.vec(1, T.f32), p_lds, [fx.Index(pbase + fx.Int32(j))] + )[0] + p_f16 = arith.truncf(_FX_KV.ir_type, arith.unwrap(fx.Float32(pf_j))) + p_frag = vector.insert( + p_f16, p_frag, static_position=[j], dynamic_position=[] + ) + p_frags.append(p_frag) + + # Prefetch V for both halves across all _PV_GROUPS D-groups. + v_pf = [] # [half][g][j] half=0..1 (32 toks each), g, j=0..7 + for half in range_constexpr(2): + vhalf = [] + for g in range_constexpr(_PV_GROUPS): + gvals = [] + for j in range_constexpr(8): + tok_j = ( + tile_start + + fx.Int32(half * (TILE_N // 2)) + + k_grp * fx.Int32(8) + + fx.Int32(j) + ) + d_out = fx.Int32(g * MFMA_N) + tok_qk + v_off = kv_base + tok_j * stride_km + d_out + v_val = buffer_ops.buffer_load( + v_rsrc, v_off, vec_width=1, dtype=_FX_KV + ) + gvals.append(arith.unwrap(_FX_KV(v_val))) + vhalf.append(gvals) + v_pf.append(vhalf) + + rocdl.sched_barrier(0) + + rescale_raw = arith.unwrap(rescale) + new_pv_scalars = [] + for g in range_constexpr(_PV_GROUPS): + c_acc = zero_v4 + for e in range_constexpr(4): + c_acc = vector.insert( + arith.mulf(pv_scalars[g * 4 + e], rescale_raw), + c_acc, + static_position=[e], + dynamic_position=[], + ) + + for half in range_constexpr(2): + v_frag = zero_v8h + for j in range_constexpr(8): + v_frag = vector.insert( + v_pf[half][g][j], + v_frag, + static_position=[j], + dynamic_position=[], + ) + c_acc = _mfma( + T.vec(4, T.f32), [v_frag, p_frags[half], c_acc, 0, 0, 0] + ) + + for e in range_constexpr(4): + new_pv_scalars.append( + vector.extract(c_acc, static_position=[e], dynamic_position=[]) + ) + + pv_scalars = new_pv_scalars + state_out = [arith.unwrap(new_max), arith.unwrap(new_sum)] + list( + pv_scalars + ) + results = yield state_out + + final_max = fx.Float32(results[0]) + final_sum = fx.Float32(results[1]) + final_pv_sc = [results[2 + i] for i in range(_N_PV)] + + safe_sum = fx.Float32( + arith.select( + arith.unwrap(final_sum) > c_zero, arith.unwrap(final_sum), c_one + ) + ) + inv_sum = rcp_f32(safe_sum) + out_base = b_idx * stride_qb + g_idx * stride_qg + hq_abs * stride_qh + + if const_expr(_SPLIT): + _pm_base = ( + b_idx * (num_g * split_total * num_hq) + + g_idx * (split_total * num_hq) + + split_idx * num_hq + + hq_abs + ) + _po_base = _pm_base * fx.Int32(_HEAD) + + # Only tok_qk=0 lanes write; d_out = g*MFMA_N + k_grp*4 + elem (C layout). + if tok_qk == fx.Int32(0): + if hq_abs < num_hq: + if const_expr(_SPLIT): + for g in range_constexpr(_PV_GROUPS): + for e in range_constexpr(4): + d_out = ( + fx.Int32(g * MFMA_N) + k_grp * fx.Int32(4) + fx.Int32(e) + ) + buffer_ops.buffer_store( + final_pv_sc[g * 4 + e], out_rsrc, _po_base + d_out + ) + else: + for g in range_constexpr(_PV_GROUPS): + for e in range_constexpr(4): + d_out = ( + fx.Int32(g * MFMA_N) + k_grp * fx.Int32(4) + fx.Int32(e) + ) + out_val = _FX_OUT( + arith.unwrap( + fx.Float32( + arith.mulf( + final_pv_sc[g * 4 + e], + arith.unwrap(inv_sum), + ) + ) + ) + ) + buffer_ops.buffer_store( + arith.unwrap(out_val), out_rsrc, out_base + d_out + ) + + if const_expr(_SPLIT): + if lane == fx.Int32(0): + if hq_abs < num_hq: + buffer_ops.buffer_store(arith.unwrap(final_max), pm_rsrc, _pm_base) + buffer_ops.buffer_store(arith.unwrap(final_sum), ps_rsrc, _pm_base) + + return pa_decode_generic_kernel, alloc + + +@functools.lru_cache(maxsize=256) +def _make_generic_jit_launcher( + head_size: int, + kv_dtype_str: str, + out_dtype_str: str, + split_k: int, +) -> Any: # pyre-ignore[3] + kernel, _alloc = compile_pa_decode_generic( + head_size=head_size, + kv_dtype_str=kv_dtype_str, + output_dtype_str=out_dtype_str, + split_k=split_k, + ) + + @flyc.jit + def _launcher( + out_ptr: fx.Tensor, + pm_ptr: fx.Tensor, + ps_ptr: fx.Tensor, + q_ptr: fx.Tensor, + k_ptr: fx.Tensor, + v_ptr: fx.Tensor, + seq_ptr: fx.Tensor, + stride_qb: fx.Int32, + stride_qg: fx.Int32, + stride_qh: fx.Int32, + stride_kb: fx.Int32, + stride_km: fx.Int32, + stride_kg: fx.Int32, + stride_kh: fx.Int32, + num_hq: fx.Int32, + num_g: fx.Int32, + kv_max: fx.Int32, + num_hkv: fx.Int32, + scale: fx.Float32, + split_total: fx.Int32, + grid_x: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ) -> None: + from flydsl._mlir import ir as _ir # pyre-ignore[21] + from flydsl.compiler.kernel_function import ( # pyre-ignore[21] + CompilationContext, + ) + + _alloc.finalized = False + ctx = CompilationContext.get_current() + with _ir.InsertionPoint(ctx.gpu_module_body): + _alloc.finalize() + kernel( + out_ptr, + pm_ptr, + ps_ptr, + q_ptr, + k_ptr, + v_ptr, + seq_ptr, + stride_qb, + stride_qg, + stride_qh, + stride_kb, + stride_km, + stride_kg, + stride_kh, + num_hq, + num_g, + kv_max, + num_hkv, + scale, + split_total, + ).launch(grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + return _launcher + + +def pa_decode_generic_launch( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + seq_positions: Optional[torch.Tensor], + softmax_scale: float, + split_k: int = 0, + output_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + """Decode: mfma_f32_16x16x32 + per-warp Q-head ownership + TLOOP.""" + from mslk.flydsl.jit import run_compiled # pyre-ignore[21] + + from .pa_decode_dense import auto_split_k + + B, _, G, H_q, D = Q.shape + _, KV_MAX, _, H_kv, _ = K.shape + assert D % MFMA_K == 0, f"head_size must be multiple of {MFMA_K}" + assert K.dtype in (torch.float16, torch.bfloat16) + + if output_dtype is None: + output_dtype = Q.dtype + kv_str = {torch.float16: "f16", torch.bfloat16: "bf16"}[K.dtype] + out_str = {torch.float16: "f16", torch.bfloat16: "bf16", torch.float32: "f32"}[ + output_dtype + ] + + if seq_positions is None: + seq_positions = torch.full((B,), KV_MAX, dtype=torch.int32, device=Q.device) + elif seq_positions.dtype != torch.int32: + seq_positions = seq_positions.to(torch.int32) + + if split_k == 0: + split_k = auto_split_k(B, G, H_q, KV_MAX) + + hq_blocks = (H_q + NUM_WARPS - 1) // NUM_WARPS + out = torch.empty((B, 1, G, H_q, D), dtype=output_dtype, device=Q.device) + sq = Q.stride() + sk2 = K.stride() + dev = Q.device + # Thread the live stream into .launch so the kernel is captured under CUDA graphs + # (a default-stream launch would capture empty). + stream = torch.cuda.current_stream() + + if split_k == 1: + dummy = torch.empty(0, dtype=torch.float32, device=dev) + launcher = _make_generic_jit_launcher(D, kv_str, out_str, 1) + grid_x = B * G * hq_blocks + run_compiled( + launcher, + out, + dummy, + dummy, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + softmax_scale, + split_k, + grid_x, + stream, + ) + else: + po = torch.empty((B, G, split_k, H_q, D), dtype=torch.float32, device=dev) + pm = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + ps = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + launcher = _make_generic_jit_launcher(D, kv_str, "f32", split_k) + grid_x = B * G * hq_blocks * split_k + run_compiled( + launcher, + po, + pm, + ps, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + softmax_scale, + split_k, + grid_x, + stream, + ) + out_view = out.squeeze(1) + pa_decode_reduce(po, pm, ps, out_view, stream=stream) + + return out diff --git a/mslk/attention/fmha/flydsl/pa_decode_gfx950.py b/mslk/attention/fmha/flydsl/pa_decode_gfx950.py new file mode 100644 index 00000000..a89e3af8 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_gfx950.py @@ -0,0 +1,667 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL decode (gfx950) — head-packed MFMA + double-buffered wide V load. + +Packs up to 16 query heads sharing a KV head onto the MFMA M-dim. One CTA = one +warp = one KV head's whole GQA group. + QK: A=Q[head=M(16), k=head_dim], B=K[tok=N(16), k=head_dim] -> C[head, tok]. + Softmax: per-head max/sum reduce over low 4 lane bits via dpp_xor(1,2,4,8). + PV: A=P[head=M, tok=K(32)], B=V[tok=K, d=N(16)] -> C[head, d]. +MFMA reg<->matrix layout: 16x16x32 lane l reg e -> C[m=(l//16)*4+e, n=l%16]. + +V staged into LDS in [dpass][tok][16] transpose layout via wide vec8 loads, read +back with ds_read_tr16_b64. V HBM loads issued EARLY so latency overlaps QK+softmax +(intra-tile software pipeline). + +gfx950 only; GQA ratio in [1,16] (else falls back to pa_decode_generic). Split-K +via pa_decode_reduce. +""" + +from __future__ import annotations + +import functools +from typing import Any + +import flydsl.compiler as flyc # pyre-ignore[21] +import flydsl.expr as fx # pyre-ignore[21] +import torch +from flydsl._mlir.dialects import llvm as _llvm # pyre-ignore[21] +from flydsl.expr import ( # pyre-ignore[21] + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr # pyre-ignore[21] + +from .pa_decode_reduce import pa_decode_reduce +from .utils import ( + dpp_xor_f32, + exp2_f32 as _exp2_fast, + maxnumf as _mxf, + rcp_f32, + smem_bytes, + WARP_SIZE, +) + +MFMA_M = 16 # heads packed on the QK MFMA M-axis +MFMA_N = 16 # tokens per QK sub-tile (N-axis) +MFMA_K_QK = 32 # QK MFMA K-dim (head-dim elements per call) +TILE_N = 32 # tokens per streaming tile (= PV MFMA K-dim) +N_SUBTILE = TILE_N // MFMA_N # 2 QK sub-tiles per tile +BLOCK = WARP_SIZE # one warp per CTA + +_FX_DTYPE = {"f32": fx.Float32, "f16": fx.Float16, "bf16": fx.BFloat16} +LOG2E: float = 1.4426950408889634 + + +@functools.lru_cache(maxsize=256) +def compile_pa_decode_gfx950( + *, + head_size: int, + kv_dtype_str: str, + output_dtype_str: str, + split_k: int = 1, + arch: str = "", +) -> Any: # pyre-ignore[3] + if not arch: + arch = get_rocm_arch() + assert head_size % MFMA_K_QK == 0 + assert head_size % MFMA_N == 0 + assert kv_dtype_str in ("f16", "bf16") + assert arch.startswith("gfx950"), f"pa_decode_gfx950 requires gfx950, got {arch}" + + _HEAD = head_size + _SK = split_k + _SPLIT = _SK > 1 + _FX_KV = _FX_DTYPE[kv_dtype_str] + _FX_OUT = _FX_DTYPE[output_dtype_str] + _QK_GRP = _HEAD // MFMA_K_QK # head-dim groups for QK (4 at D=128) + _DN = _HEAD // MFMA_N # d-passes for PV (8 at D=128) + _mfma = ( + rocdl.mfma_f32_16x16x32_bf16 + if kv_dtype_str == "bf16" + else rocdl.mfma_f32_16x16x32_f16 + ) + + # LDS: P[MFMA_M, TILE_N] f32 + double-buffered V ([dpass][tok][16] transpose tiles). + _NUM_DMA_V = (TILE_N * _HEAD // 8) // WARP_SIZE # 16B (8 f16) chunks / 64 lanes + _P_LDS = MFMA_M * TILE_N # f32, P redistribution + _V_LDS = TILE_N * _HEAD # f16, one V tile (transpose layout) + _P_BYTES = _P_LDS * 4 + _V_BYTES = _V_LDS * 2 # per buffer + _LDS_TOTAL = _P_BYTES + 2 * _V_BYTES # double-buffered V + cap = smem_bytes(arch) + if _LDS_TOTAL > cap: + raise ValueError(f"LDS {_LDS_TOTAL}B > {arch!r} cap {cap}B") + + alloc = SmemAllocator( + None, + arch=arch, + global_sym_name=f"pa_gfx950_h{_HEAD}_{kv_dtype_str}_sk{_SK}", + ) + alloc.ptr = _LDS_TOTAL + + @flyc.kernel(known_block_size=(BLOCK, 1, 1)) + def pa_decode_gfx950_kernel( + out_ptr: fx.Tensor, + partial_max_ptr: fx.Tensor, + partial_sum_ptr: fx.Tensor, + q_ptr: fx.Tensor, + k_ptr: fx.Tensor, + v_ptr: fx.Tensor, + seq_ptr: fx.Tensor, + stride_qb: fx.Int32, + stride_qg: fx.Int32, + stride_qh: fx.Int32, + stride_kb: fx.Int32, + stride_km: fx.Int32, + stride_kg: fx.Int32, + stride_kh: fx.Int32, + num_hq: fx.Int32, + num_g: fx.Int32, + kv_max: fx.Int32, + num_hkv: fx.Int32, + ratio: fx.Int32, + softmax_scale: fx.Float32, + split_total: fx.Int32, + ) -> None: + lane = gpu.thread_idx.x + tok_lane = lane % fx.Int32(MFMA_N) # 0..15 (N index / token within sub-tile) + grp = lane // fx.Int32(MFMA_N) # 0..3 (M-group / k-sub-group) + + # Grid: flat -> (split_idx, kv_head, g, b). One CTA per (b,g,kv_head[,split]). + flat = fx.Int32(gpu.block_idx.x) + if const_expr(_SPLIT): + split_idx = flat % split_total + rest = flat // split_total + else: + split_idx = fx.Int32(0) + rest = flat + hkv_abs = rest % num_hkv + rest2 = rest // num_hkv + g_idx = rest2 % num_g + b_idx = rest2 // num_g + hq_base = hkv_abs * ratio # first query head sharing this KV head + + c_zero = arith.constant(0.0, type=T.f32) + c_one = arith.constant(1.0, type=T.f32) + c_neginf = arith.constant(float("-inf"), type=T.f32) + zero_v4 = arith.constant_vector(0.0, T.vec(4, T.f32)) + zero_v8h = arith.constant_vector(0.0, T.vec(8, _FX_KV.ir_type)) + + q_rsrc = buffer_ops.create_buffer_resource(q_ptr, max_size=True) + k_rsrc = buffer_ops.create_buffer_resource(k_ptr, max_size=True) + v_rsrc = buffer_ops.create_buffer_resource(v_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True) + pm_rsrc = buffer_ops.create_buffer_resource(partial_max_ptr, max_size=True) + ps_rsrc = buffer_ops.create_buffer_resource(partial_sum_ptr, max_size=True) + seq_rsrc = buffer_ops.create_buffer_resource(seq_ptr, max_size=True) + + seq_len = buffer_ops.buffer_load(seq_rsrc, b_idx, vec_width=1, dtype=T.i32) + t_full = arith.select(seq_len > fx.Int32(0), seq_len, kv_max) + if const_expr(_SPLIT): + chunk = (t_full + split_total - fx.Int32(1)) // split_total + t_start = split_idx * chunk + t_end_raw = (split_idx + fx.Int32(1)) * chunk + t_end = arith.select(t_end_raw < t_full, t_end_raw, t_full) + else: + t_start = fx.Int32(0) + t_end = t_full + + smem = alloc.get_base() + p_lds = SmemPtr(smem, 0, T.f32, shape=(_P_LDS,)).get() + lds_base = buffer_ops.extract_base_index(smem, address_space=3) # f16 elem base + v_lds_f16 = lds_base + fx.Index( + _P_BYTES // 2 + ) # V tile (transpose) after P region + + # ── Pre-load Q (loop-invariant) ── + # A-frag: lane l -> Q[head=tok_lane, k=grp*8+0..7]. head on M = tok_lane + # (0..15); only heads < ratio meaningful. + q_head = tok_lane + q_base = b_idx * stride_qb + g_idx * stride_qg + (hq_base + q_head) * stride_qh + q_frags = [] + for g in range_constexpr(_QK_GRP): + q_off = q_base + fx.Int32(g * MFMA_K_QK) + grp * fx.Int32(8) + q_frags.append( + buffer_ops.buffer_load(q_rsrc, q_off, vec_width=8, dtype=_FX_KV) + ) + + kv_base = b_idx * stride_kb + g_idx * stride_kg + hkv_abs * stride_kh + + # Loop-carried state: per-head (reg e over 0..3) running max, running sum, + # and PV accumulator (_DN d-passes x 4 regs). + _N_ACC = _DN * 4 + _init = [c_neginf] * 4 + [c_zero] * 4 + [c_zero] * _N_ACC + + for _tile_i, state in range( + fx.Index(t_start), fx.Index(t_end), arith.index(TILE_N), init=_init + ): + rmax = [fx.Float32(state[i]) for i in range(4)] + rsum = [fx.Float32(state[4 + i]) for i in range(4)] + acc = [fx.Float32(state[8 + i]) for i in range(_N_ACC)] + tile_start = fx.Int32(arith.index_cast(T.i32, _tile_i)) + + # ── Issue V HBM loads EARLY (into regs) so latency overlaps the + # QK+softmax below; LDS transpose stores + barrier happen just before PV. + _v8s = [] + _v8dst = [] + for _r in range_constexpr(_NUM_DMA_V): + _lin = lane + fx.Int32(_r * WARP_SIZE) + _tok = _lin % fx.Int32(TILE_N) + _rest = _lin // fx.Int32(TILE_N) # dpass*2 + half + _dp = _rest // fx.Int32(2) + _half = _rest % fx.Int32(2) + _col = _dp * fx.Int32(16) + _half * fx.Int32(8) + _v8s.append( + buffer_ops.buffer_load( + v_rsrc, + kv_base + (tile_start + _tok) * stride_km + _col, + vec_width=8, + dtype=_FX_KV, + ) + ) + _dst = ( + v_lds_f16 + + fx.Index(_dp) * fx.Index(TILE_N * 16) + + fx.Index(_tok) * fx.Index(16) + + fx.Index(_half) * fx.Index(8) + ) + _v8dst.append(_dst) + + # ── QK: N_SUBTILE sub-tiles of 16 tokens ── + # qk[st] reg e -> score[head=grp*4+e, tok=st*16+tok_lane] + qk_st = [] + for st in range_constexpr(N_SUBTILE): + acc_qk = zero_v4 + for g in range_constexpr(_QK_GRP): + k_tok = tile_start + fx.Int32(st * MFMA_N) + tok_lane + k_off = ( + kv_base + + k_tok * stride_km + + fx.Int32(g * MFMA_K_QK) + + grp * fx.Int32(8) + ) + k8 = buffer_ops.buffer_load( + k_rsrc, k_off, vec_width=8, dtype=_FX_KV + ) + acc_qk = _mfma(T.vec(4, T.f32), [q_frags[g], k8, acc_qk, 0, 0, 0]) + qk_st.append(acc_qk) + + # ── Online softmax, per head (reg e); tile max via dpp over tok_lane ── + new_max = [] + alpha = [] + for e in range_constexpr(4): + loc = fx.Float32(c_neginf) + for st in range_constexpr(N_SUBTILE): + s = fx.Float32( + vector.extract( + qk_st[st], static_position=[e], dynamic_position=[] + ) + ) + s = fx.Float32( + arith.mulf(arith.unwrap(s), arith.unwrap(softmax_scale)) + ) + # mask out-of-range tokens + tok_abs = tile_start + fx.Int32(st * MFMA_N) + tok_lane + ok = tok_abs < t_end + s = fx.Float32( + arith.select(arith.unwrap(ok), arith.unwrap(s), c_neginf) + ) + loc = _mxf(loc, s) + qk_st[st] = vector.insert( + arith.unwrap(s), + qk_st[st], + static_position=[e], + dynamic_position=[], + ) + for sh in (1, 2, 4, 8): + loc = _mxf(loc, dpp_xor_f32(loc, sh)) + nm = _mxf(rmax[e], loc) + new_max.append(nm) + a = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf(arith.unwrap(rmax[e]), arith.unwrap(nm)), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + alpha.append(a) + + # P = exp2((score - new_max)*log2e); write to LDS[head, tok]; accumulate sum. + tile_sum = [fx.Float32(c_zero) for _ in range(4)] + for e in range_constexpr(4): + head = grp * fx.Int32(4) + fx.Int32(e) + for st in range_constexpr(N_SUBTILE): + s = fx.Float32( + vector.extract( + qk_st[st], static_position=[e], dynamic_position=[] + ) + ) + p = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf(arith.unwrap(s), arith.unwrap(new_max[e])), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + # masked lanes gave s=-inf -> p=0 + p = fx.Float32( + arith.select( + arith.unwrap(new_max[e]) > c_neginf, arith.unwrap(p), c_zero + ) + ) + tile_sum[e] = fx.Float32( + arith.addf(arith.unwrap(tile_sum[e]), arith.unwrap(p)) + ) + tok = fx.Int32(st * MFMA_N) + tok_lane + vector.store( + fx.Vector.from_elements([arith.unwrap(p)], dtype=fx.Float32), + p_lds, + [fx.Index(head * fx.Int32(TILE_N) + tok)], + ) + for e in range_constexpr(4): + for sh in (1, 2, 4, 8): + tile_sum[e] = fx.Float32( + arith.addf( + arith.unwrap(tile_sum[e]), + arith.unwrap(dpp_xor_f32(tile_sum[e], sh)), + ) + ) + rsum[e] = fx.Float32( + arith.addf( + arith.mulf(arith.unwrap(alpha[e]), arith.unwrap(rsum[e])), + arith.unwrap(tile_sum[e]), + ) + ) + rmax[e] = new_max[e] + + # Write the (already-loaded) V vec8s into the LDS transpose layout; the + # barrier below covers both P writes and these V writes before PV. + for _r in range_constexpr(_NUM_DMA_V): + _sp = buffer_ops.create_llvm_ptr( + fx.Int64(_v8dst[_r] * fx.Index(2)), address_space=3 + ) + _llvm.StoreOp(_v8s[_r], _sp, alignment=16) + + gpu.barrier() + + # ── PV: A=P[head,tok] (LDS), B=V[tok,d] -> C[head,d]; rescale acc by alpha ── + for dpass in range_constexpr(_DN): + for e in range_constexpr(4): + acc[dpass * 4 + e] = fx.Float32( + arith.mulf( + arith.unwrap(acc[dpass * 4 + e]), arith.unwrap(alpha[e]) + ) + ) + # A-frag P: lane l -> P[head = tok_lane, tok = grp*8 + 0..7] + p_head = tok_lane + p_vals = [] + for j in range_constexpr(8): + pv = fx.Vector.load( + T.vec(1, T.f32), + p_lds, + [ + fx.Index( + p_head * fx.Int32(TILE_N) + grp * fx.Int32(8) + fx.Int32(j) + ) + ], + )[0] + p_vals.append( + arith.truncf(_FX_KV.ir_type, arith.unwrap(fx.Float32(pv))) + ) + p_frag = zero_v8h + for j in range_constexpr(8): + p_frag = vector.insert( + p_vals[j], p_frag, static_position=[j], dynamic_position=[] + ) + + _v4h = T.vec(4, _FX_KV.ir_type) + for dpass in range_constexpr(_DN): + # B-frag V via two ds_read_tr16_b64 (128-bit HW transpose): group grp + # owns toks grp*8..+7 -> V[tok, d=dpass*16+tok_lane] (2 wide reads vs 8). + _GB = fx.Int32(dpass * (TILE_N * 16)) + (grp * fx.Int32(8)) * fx.Int32( + 16 + ) + _off_lo = v_lds_f16 + fx.Index(_GB) + (fx.Index(tok_lane)) * fx.Index(4) + _vlo = rocdl.ds_read_tr16_b64( + _v4h, + buffer_ops.create_llvm_ptr( + fx.Int64(_off_lo * fx.Index(2)), address_space=3 + ), + ).result + _off_hi = _off_lo + fx.Index(4 * 16) + _vhi = rocdl.ds_read_tr16_b64( + _v4h, + buffer_ops.create_llvm_ptr( + fx.Int64(_off_hi * fx.Index(2)), address_space=3 + ), + ).result + v_frag = vector.shuffle(_vlo, _vhi, [0, 1, 2, 3, 4, 5, 6, 7]) + c_in = zero_v4 + for e in range_constexpr(4): + c_in = vector.insert( + arith.unwrap(acc[dpass * 4 + e]), + c_in, + static_position=[e], + dynamic_position=[], + ) + c_out = _mfma(T.vec(4, T.f32), [p_frag, v_frag, c_in, 0, 0, 0]) + for e in range_constexpr(4): + acc[dpass * 4 + e] = fx.Float32( + vector.extract(c_out, static_position=[e], dynamic_position=[]) + ) + + gpu.barrier() # P_LDS reused next tile + + state_out = ( + [arith.unwrap(rmax[i]) for i in range(4)] + + [arith.unwrap(rsum[i]) for i in range(4)] + + [arith.unwrap(acc[i]) for i in range(_N_ACC)] + ) + results = yield state_out + + f_max = [fx.Float32(results[i]) for i in range(4)] + f_sum = [fx.Float32(results[4 + i]) for i in range(4)] + f_acc = [fx.Float32(results[8 + i]) for i in range(_N_ACC)] + + # ── Epilogue: normalize + store per head ── + # head=grp*4+e; d=dpass*16+tok_lane always < _HEAD, so no d guard. + for e in range_constexpr(4): + head = grp * fx.Int32(4) + fx.Int32(e) + head_abs = hq_base + head + safe_sum = fx.Float32( + arith.select( + arith.unwrap(f_sum[e]) > c_zero, arith.unwrap(f_sum[e]), c_one + ) + ) + inv = rcp_f32(safe_sum) + if const_expr(_SPLIT): + _pm_base = ( + b_idx * (num_g * split_total * num_hq) + + g_idx * (split_total * num_hq) + + split_idx * num_hq + + head_abs + ) + _po_base = _pm_base * fx.Int32(_HEAD) + if (head < ratio) & (head_abs < num_hq): + for dpass in range_constexpr(_DN): + d = fx.Int32(dpass * MFMA_N) + tok_lane + buffer_ops.buffer_store( + arith.unwrap(f_acc[dpass * 4 + e]), out_rsrc, _po_base + d + ) + if tok_lane == fx.Int32(0): + buffer_ops.buffer_store( + arith.unwrap(f_max[e]), pm_rsrc, _pm_base + ) + buffer_ops.buffer_store( + arith.unwrap(f_sum[e]), ps_rsrc, _pm_base + ) + else: + out_base = b_idx * stride_qb + g_idx * stride_qg + head_abs * stride_qh + inv_raw = arith.unwrap(inv) + if (head < ratio) & (head_abs < num_hq): + for dpass in range_constexpr(_DN): + d = fx.Int32(dpass * MFMA_N) + tok_lane + val = fx.Float32( + arith.mulf(arith.unwrap(f_acc[dpass * 4 + e]), inv_raw) + ) + out_val = _FX_OUT(arith.unwrap(val)) + buffer_ops.buffer_store( + arith.unwrap(out_val), out_rsrc, out_base + d + ) + + return pa_decode_gfx950_kernel, alloc + + +@functools.lru_cache(maxsize=256) +def _make_gfx950_jit_launcher(head_size, kv_dtype_str, out_dtype_str, split_k): + kernel, _alloc = compile_pa_decode_gfx950( + head_size=head_size, + kv_dtype_str=kv_dtype_str, + output_dtype_str=out_dtype_str, + split_k=split_k, + ) + + @flyc.jit + def _launcher( + out_ptr, + pm_ptr, + ps_ptr, + q_ptr, + k_ptr, + v_ptr, + seq_ptr, + stride_qb, + stride_qg, + stride_qh, + stride_kb, + stride_km, + stride_kg, + stride_kh, + num_hq, + num_g, + kv_max, + num_hkv, + ratio, + scale, + split_total, + grid_x, + stream: fx.Stream = fx.Stream(None), + ): + from flydsl._mlir import ir as _ir + from flydsl.compiler.kernel_function import CompilationContext + + _alloc.finalized = False + ctx = CompilationContext.get_current() + with _ir.InsertionPoint(ctx.gpu_module_body): + _alloc.finalize() + kernel( + out_ptr, + pm_ptr, + ps_ptr, + q_ptr, + k_ptr, + v_ptr, + seq_ptr, + stride_qb, + stride_qg, + stride_qh, + stride_kb, + stride_km, + stride_kg, + stride_kh, + num_hq, + num_g, + kv_max, + num_hkv, + ratio, + scale, + split_total, + ).launch(grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + return _launcher + + +def pa_decode_gfx950_launch( + Q, K, V, seq_positions, softmax_scale, split_k=0, output_dtype=None +): + """Head-packed MFMA decode. One CTA per KV head packs its GQA group onto + the MFMA M-axis. Falls back to the generic kernel for ratio>16 or non-gfx950.""" + from mslk.flydsl.jit import run_compiled + + from .pa_decode_dense import auto_split_k_hp + + B, _, G, H_q, D = Q.shape + _, KV_MAX, _, H_kv, _ = K.shape + ratio = H_q // H_kv if H_kv > 0 else 0 + ok = ( + H_kv > 0 + and H_q % H_kv == 0 + and 1 <= ratio <= MFMA_M + and get_rocm_arch().startswith("gfx950") + and K.dtype in (torch.float16, torch.bfloat16) + and D % MFMA_K_QK == 0 + ) + if not ok: + from .pa_decode_generic import pa_decode_generic_launch + + return pa_decode_generic_launch( + Q, K, V, seq_positions, softmax_scale, split_k, output_dtype + ) + if output_dtype is None: + output_dtype = Q.dtype + kv_str = {torch.float16: "f16", torch.bfloat16: "bf16"}[K.dtype] + out_str = {torch.float16: "f16", torch.bfloat16: "bf16", torch.float32: "f32"}[ + output_dtype + ] + if seq_positions is None: + seq_positions = torch.full((B,), KV_MAX, dtype=torch.int32, device=Q.device) + elif seq_positions.dtype != torch.int32: + seq_positions = seq_positions.to(torch.int32) + if split_k == 0: + split_k = auto_split_k_hp(B, G, H_q, H_kv, KV_MAX) + out = torch.empty((B, 1, G, H_q, D), dtype=output_dtype, device=Q.device) + sq = Q.stride() + sk2 = K.stride() + dev = Q.device + n_cta_base = B * G * H_kv + # Thread the live stream into .launch so the kernel is captured under CUDA graphs + # (a default-stream launch would capture empty). + stream = torch.cuda.current_stream() + if split_k == 1: + dummy = torch.empty(0, dtype=torch.float32, device=dev) + launcher = _make_gfx950_jit_launcher(D, kv_str, out_str, 1) + run_compiled( + launcher, + out, + dummy, + dummy, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + ratio, + softmax_scale, + split_k, + n_cta_base, + stream, + ) + else: + po = torch.empty((B, G, split_k, H_q, D), dtype=torch.float32, device=dev) + pm = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + ps = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + launcher = _make_gfx950_jit_launcher(D, kv_str, "f32", split_k) + run_compiled( + launcher, + po, + pm, + ps, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + ratio, + softmax_scale, + split_k, + n_cta_base * split_k, + stream, + ) + pa_decode_reduce(po, pm, ps, out.squeeze(1), stream=stream) + return out diff --git a/mslk/attention/fmha/flydsl/pa_decode_gfx950_coop.py b/mslk/attention/fmha/flydsl/pa_decode_gfx950_coop.py new file mode 100644 index 00000000..08dc4f53 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_gfx950_coop.py @@ -0,0 +1,730 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL decode (gfx950 cooperative-DMA) — ds_read_tr16_b64 HW transpose. + +Per-head coop-DMA fallback for GQA ratios that can't head-pack. Uses ds_read_tr16_b64 +(gfx950+ HW LDS transpose) for PV to cut V LDS reads 8x. + +PV MFMA mfma_f32_32x32x16_f16, A=V_T (from ds_read_tr16_b64), B=P (broadcast): + A[m=d_sub, k=tok] = V[tok, d=dc*32+d_sub]; C[m=d_sub, n=*] = PV[d=dc*32+d_sub]. + +Lane decomposition for ds_read_tr16_b64: + lane_div_32 = lane//32 -> tok half within pks step + tr_k_group = (lane%16)//4 -> K-row (tok) offset within 4-row group + tr_col_sub = lane%4 -> 4-column (d) sub-group + tr_col_half = (lane%32)//16-> first/second 16-d half of DC chunk + +V LDS is linear row-major (required by ds_read_tr16_b64). gfx950 only. +""" + +from __future__ import annotations + +import functools +from typing import Any + +import flydsl.compiler as flyc # pyre-ignore[21] +import flydsl.expr as fx # pyre-ignore[21] +import torch +from flydsl._mlir.dialects import llvm as _llvm # pyre-ignore[21] +from flydsl.expr import ( # pyre-ignore[21] + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr # pyre-ignore[21] + +from .pa_decode_reduce import pa_decode_reduce +from .utils import ( + dpp_xor_f32, + exp2_f32 as _exp2_fast, + maxnumf as _mxf, + rcp_f32, + smem_bytes, + WARP_SIZE, +) + +NUM_WARPS = 4 +MFMA_N = 16 # QK MFMA sub-tile (tokens per group, mfma_f32_16x16x32_f16) +MFMA_K_QK = 32 # QK MFMA K-dim +TLOOP = NUM_WARPS # 4 sub-tiles per tile +TILE_N = TLOOP * MFMA_N # 64 tokens per tile +BLOCK = NUM_WARPS * WARP_SIZE # 256 + +DMA_BYTES = 16 # bytes per lane per DMA call (raw_ptr_buffer_load_lds) + +# PV: mfma_f32_32x32x16_f16 with ds_read_tr16_b64 +DC_CHUNK = 32 # d-values per DC pass (MFMA M=32); _D_CHUNKS = HEAD//DC_CHUNK +PV_K_STEP = 16 # tokens per pks step (MFMA K=16) +K_SUB_N = 32 # half TILE_N (lo vs hi token groups) +PV_K_STEPS = TILE_N // PV_K_STEP # 4 steps: pks=0..3 + +_FX_DTYPE = {"f32": fx.Float32, "f16": fx.Float16, "bf16": fx.BFloat16} +LOG2E: float = 1.4426950408889634 + + +@functools.lru_cache(maxsize=256) +def compile_pa_decode_gfx950_coop( + *, + head_size: int, + kv_dtype_str: str, + output_dtype_str: str, + split_k: int = 1, + arch: str = "", +) -> Any: # pyre-ignore[3] + if not arch: + arch = get_rocm_arch() + assert head_size % MFMA_K_QK == 0 + assert head_size % DC_CHUNK == 0 + assert kv_dtype_str in ("f16", "bf16") + assert arch.startswith("gfx950"), ( + f"pa_decode_gfx950_coop requires gfx950 (ds_read_tr16_b64), got {arch}" + ) + + _HEAD = head_size + _SK = split_k + _SPLIT = _SK > 1 + _FX_KV = _FX_DTYPE[kv_dtype_str] + _FX_OUT = _FX_DTYPE[output_dtype_str] + # MFMA intrinsic must match KV operand dtype (mismatch fails MLIR verification). + _mfma_qk = ( + rocdl.mfma_f32_16x16x32_bf16 + if kv_dtype_str == "bf16" + else rocdl.mfma_f32_16x16x32_f16 + ) + _mfma_pv = ( + rocdl.mfma_f32_32x32x16_bf16 + if kv_dtype_str == "bf16" + else rocdl.mfma_f32_32x32x16_f16 + ) + _QK_GROUPS = _HEAD // MFMA_K_QK # 4 for head=128 + _D_CHUNKS = _HEAD // DC_CHUNK # head-dim / 32 (=4 for head=128, 8 for head=256) + # PV accumulator: _D_CHUNKS * 16 scalars per lane (v16f32 per DC chunk) + _N_PV = _D_CHUNKS * 16 + + # LDS: K_LDS TILE_N x HEAD f16 (XOR swizzle) + V_LDS TILE_N x HEAD f16 + # (row-major, no swizzle) + P_LDS NUM_WARPS x TILE_N f32. + _K_LDS_F16 = TILE_N * _HEAD + _V_LDS_F16 = TILE_N * _HEAD + _P_LDS_F32 = NUM_WARPS * TILE_N + _LDS_TOTAL = (_K_LDS_F16 + _V_LDS_F16) * 2 + _P_LDS_F32 * 4 + + cap = smem_bytes(arch) + if _LDS_TOTAL > cap: + raise ValueError(f"LDS {_LDS_TOTAL}B > {arch!r} cap {cap}B") + + alloc = SmemAllocator( + None, + arch=arch, + global_sym_name=f"pa_gfx950_coop_h{_HEAD}_{kv_dtype_str}_nw{NUM_WARPS}_sk{_SK}", + ) + alloc.ptr = _LDS_TOTAL + + _DMA_BATCH = BLOCK * DMA_BYTES + _KV_TILE_BYTES = TILE_N * _HEAD * 2 + _NUM_DMA_KV = _KV_TILE_BYTES // _DMA_BATCH # 4 rounds + _LANES_PER_ROW = _HEAD * 2 // DMA_BYTES # 16 + _ROWS_PER_ROUND = _DMA_BATCH // (_HEAD * 2) # 16 + + # V LDS stride (row-major; ds_read_tr16_b64 needs no padding) + _V_STRIDE = _HEAD # f16 per row (tok) + + @flyc.kernel(known_block_size=(BLOCK, 1, 1)) + def pa_decode_gfx950_coop_kernel( + out_ptr: fx.Tensor, + partial_max_ptr: fx.Tensor, + partial_sum_ptr: fx.Tensor, + q_ptr: fx.Tensor, + k_ptr: fx.Tensor, + v_ptr: fx.Tensor, + seq_ptr: fx.Tensor, + stride_qb: fx.Int32, + stride_qg: fx.Int32, + stride_qh: fx.Int32, + stride_kb: fx.Int32, + stride_km: fx.Int32, + stride_kg: fx.Int32, + stride_kh: fx.Int32, + num_hq: fx.Int32, + num_g: fx.Int32, + kv_max: fx.Int32, + num_hkv: fx.Int32, + softmax_scale: fx.Float32, + split_total: fx.Int32, + ) -> None: + tid = gpu.thread_idx.x + warp_id = tid >> fx.Int32(6) + lane = tid & fx.Int32(63) + + # QK lane decomposition (mfma_f32_16x16x32_f16) + tok_qk = lane & fx.Int32(MFMA_N - 1) + k_grp = lane >> fx.Int32(4) + + # ds_read_tr16_b64 lane decomposition (PV mfma_f32_32x32x16_f16) + lane_div_32 = lane >> fx.Int32(5) # 0 or 1 + tr_k_group = (lane & fx.Int32(15)) >> fx.Int32(2) # (lane%16)//4: 0..3 + tr_col_sub = lane & fx.Int32(3) # lane%4: 0..3 + tr_col_half = (lane & fx.Int32(31)) >> fx.Int32(4) # (lane%32)//16: 0 or 1 + + # Grid decode + flat = fx.Int32(gpu.block_idx.x) + if const_expr(_SPLIT): + split_idx = flat % split_total + rest = flat // split_total + else: + split_idx = fx.Int32(0) + rest = flat + + n_hq_blocks = (num_hq + fx.Int32(NUM_WARPS - 1)) // fx.Int32(NUM_WARPS) + hq_block = rest % n_hq_blocks + rest2 = rest // n_hq_blocks + g_idx = rest2 % num_g + b_idx = rest2 // num_g + + hq_abs = hq_block * fx.Int32(NUM_WARPS) + warp_id + hkv_abs = hq_abs * num_hkv // num_hq + q_base = b_idx * stride_qb + g_idx * stride_qg + hq_abs * stride_qh + kv_base = b_idx * stride_kb + g_idx * stride_kg + hkv_abs * stride_kh + + c_zero = arith.constant(0.0, type=T.f32) + c_one = arith.constant(1.0, type=T.f32) + c_neginf = arith.constant(float("-inf"), type=T.f32) + zero_v4 = arith.constant_vector(0.0, T.vec(4, T.f32)) + zero_v8h = arith.constant_vector(0.0, T.vec(8, _FX_KV.ir_type)) + zero_v16 = arith.constant_vector(0.0, T.vec(16, T.f32)) + + seq_rsrc = buffer_ops.create_buffer_resource(seq_ptr, max_size=True) + q_rsrc = buffer_ops.create_buffer_resource(q_ptr, max_size=True) + k_rsrc = buffer_ops.create_buffer_resource(k_ptr, max_size=True) + v_rsrc = buffer_ops.create_buffer_resource(v_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(out_ptr, max_size=True) + pm_rsrc = buffer_ops.create_buffer_resource(partial_max_ptr, max_size=True) + ps_rsrc = buffer_ops.create_buffer_resource(partial_sum_ptr, max_size=True) + + seq_len = buffer_ops.buffer_load(seq_rsrc, b_idx, vec_width=1, dtype=T.i32) + t_full = arith.select(seq_len > fx.Int32(0), seq_len, kv_max) + if const_expr(_SPLIT): + chunk = (t_full + split_total - fx.Int32(1)) // split_total + t_start = split_idx * chunk + t_end_raw = (split_idx + fx.Int32(1)) * chunk + t_end = arith.select(t_end_raw < t_full, t_end_raw, t_full) + else: + t_start = fx.Int32(0) + t_end = t_full + + smem = alloc.get_base() + lds_base = buffer_ops.extract_base_index(smem, address_space=3) + k_lds_base_bytes = lds_base + v_lds_base_bytes = lds_base + fx.Index(_K_LDS_F16 * 2) + p_lds = SmemPtr( + smem, (_K_LDS_F16 + _V_LDS_F16) * 2, T.f32, shape=(_P_LDS_F32,) + ).get() + + _wave_dma_offset = fx.Index(warp_id * fx.Int32(WARP_SIZE * DMA_BYTES)) + _dma_size = fx.Int32(DMA_BYTES) + _dma_soff = fx.Int32(0) + _dma_off = fx.Int32(0) + _dma_aux = fx.Int32(1) + + # Pre-load Q + q_frags = [] + for g in range_constexpr(_QK_GROUPS): + q_off = q_base + fx.Int32(g * MFMA_K_QK) + k_grp * fx.Int32(8) + q_frags.append( + buffer_ops.buffer_load(q_rsrc, q_off, vec_width=8, dtype=_FX_KV) + ) + + _init_neg = arith.constant(float("-inf"), type=T.f32) + _init_zer = arith.constant(0.0, type=T.f32) + _init_state = [_init_neg, _init_zer] + [_init_zer] * _N_PV + + for _tile_i, state in range( + fx.Index(t_start), fx.Index(t_end), arith.index(TILE_N), init=_init_state + ): + running_max = fx.Float32(state[0]) + running_sum = fx.Float32(state[1]) + pv_scalars = [state[2 + i] for i in range(_N_PV)] + + tile_start = fx.Int32(arith.index_cast(T.i32, _tile_i)) + + # ── DMA K to LDS (linear, QK reads K linearly) ── + for d in range_constexpr(_NUM_DMA_KV): + row_in_tile = tid // fx.Index(_LANES_PER_ROW) + fx.Index( + d * _ROWS_PER_ROUND + ) + col_f16 = (tid % fx.Index(_LANES_PER_ROW)) * fx.Index(8) + global_row = tile_start + fx.Int32(row_in_tile) + k_voffset = ( + kv_base + global_row * stride_km + fx.Int32(col_f16) + ) * fx.Int32(2) + k_lds_rb = ( + k_lds_base_bytes + _wave_dma_offset + fx.Index(d * _DMA_BATCH) + ) + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + buffer_ops.create_llvm_ptr( + rocdl.readfirstlane(fx.Int64.ir_type, fx.Int64(k_lds_rb)), + address_space=3, + ), + _dma_size, + k_voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ── DMA V to LDS (row-major; ds_read_tr16_b64 needs linear layout) ── + for d in range_constexpr(_NUM_DMA_KV): + row_in_tile = tid // fx.Index(_LANES_PER_ROW) + fx.Index( + d * _ROWS_PER_ROUND + ) + col_f16 = (tid % fx.Index(_LANES_PER_ROW)) * fx.Index(8) + global_row = tile_start + fx.Int32(row_in_tile) + v_voffset = ( + kv_base + global_row * stride_km + fx.Int32(col_f16) + ) * fx.Int32(2) + v_lds_rb = ( + v_lds_base_bytes + _wave_dma_offset + fx.Index(d * _DMA_BATCH) + ) + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + buffer_ops.create_llvm_ptr( + rocdl.readfirstlane(fx.Int64.ir_type, fx.Int64(v_lds_rb)), + address_space=3, + ), + _dma_size, + v_voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + gpu.barrier() + + # ── QK (mfma_f32_16x16x32_f16) ── + tile_max = fx.Float32(c_neginf) + qk_scalars = [] + for td in range_constexpr(TLOOP): + k_tok = fx.Int32(td * MFMA_N) + tok_qk + k_v8s = [] + for g in range_constexpr(_QK_GROUPS): + k_col = fx.Int32(g * MFMA_K_QK) + k_grp * fx.Int32(8) + k_byte = k_lds_base_bytes + ( + fx.Index(k_tok) * fx.Index(_HEAD) + fx.Index(k_col) + ) * fx.Index(2) + k_ptr = buffer_ops.create_llvm_ptr( + fx.Int64(k_byte), address_space=3 + ) + k_v8s.append( + _llvm.LoadOp( + T.vec(8, _FX_KV.ir_type), k_ptr, alignment=16 + ).result + ) + qk_acc = zero_v4 + for g in range_constexpr(_QK_GROUPS): + qk_acc = _mfma_qk( + T.vec(4, T.f32), [q_frags[g], k_v8s[g], qk_acc, 0, 0, 0] + ) + tok_td = tile_start + fx.Int32(td * MFMA_N) + tok_qk + qk_raw = vector.extract( + qk_acc, static_position=[0], dynamic_position=[] + ) + qk_sc = arith.mulf(qk_raw, arith.unwrap(softmax_scale)) + head_ok = hq_abs < num_hq + tok_ok = tok_td < t_end + in_range = arith.andi(arith.unwrap(head_ok), arith.unwrap(tok_ok)) + qk_val = fx.Float32(arith.select(in_range, qk_sc, c_neginf)) + qk_scalars.append(qk_val) + tile_max = _mxf(tile_max, qk_val) + + for sh in (8, 4, 2, 1): + tile_max = _mxf(tile_max, dpp_xor_f32(tile_max, sh)) + + new_max = _mxf(running_max, tile_max) + rescale = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf(arith.unwrap(running_max), arith.unwrap(new_max)), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + + safe_max = fx.Float32( + arith.select( + arith.unwrap(new_max) > c_neginf, arith.unwrap(new_max), c_zero + ) + ) + intra_sum = fx.Float32(c_zero) + for td in range_constexpr(TLOOP): + tok_td = tile_start + fx.Int32(td * MFMA_N) + tok_qk + head_ok = hq_abs < num_hq + tok_ok = tok_td < t_end + in_range = arith.andi(arith.unwrap(head_ok), arith.unwrap(tok_ok)) + p_c = _exp2_fast( + fx.Float32( + arith.mulf( + arith.subf( + arith.unwrap(qk_scalars[td]), arith.unwrap(safe_max) + ), + arith.constant(LOG2E, type=T.f32), + ) + ) + ) + p_c = fx.Float32(arith.select(in_range, arith.unwrap(p_c), c_zero)) + intra_sum = fx.Float32( + arith.addf(arith.unwrap(intra_sum), arith.unwrap(p_c)) + ) + p_slot = fx.Index( + warp_id * fx.Int32(TILE_N) + fx.Int32(td * MFMA_N) + tok_qk + ) + vector.store( + fx.Vector.from_elements([arith.unwrap(p_c)], dtype=fx.Float32), + p_lds, + [p_slot], + ) + + for sh in (8, 4, 2, 1): + intra_sum = fx.Float32( + arith.addf( + arith.unwrap(intra_sum), + arith.unwrap(dpp_xor_f32(intra_sum, sh)), + ) + ) + new_sum = fx.Float32( + arith.addf( + arith.mulf(arith.unwrap(rescale), arith.unwrap(running_sum)), + arith.unwrap(intra_sum), + ) + ) + + gpu.barrier() + + # ── PV: mfma_f32_32x32x16_f16, A=V_T (ds_read_tr16_b64), B=P ── + # ds_read_tr16_b64 addressing: + # d_col = dc*32 + tr_col_half*16 + tr_col_sub*4 (f16, d-dim) + # k_row = pks*16 + ld32*4 + tr_k_group (f16, tok-dim) + # v_lo=tr16(lds_lo) -> k=0..3; v_hi=tr16(+8 toks) -> k=4..7. + + v4f16_type = T.vec(4, _FX_KV.ir_type) + + rescale_raw = arith.unwrap(rescale) + new_pv_scalars = [] + for dc in range_constexpr(_D_CHUNKS): + c_acc = zero_v16 + for e in range_constexpr(16): + c_acc = vector.insert( + arith.mulf(pv_scalars[dc * 16 + e], rescale_raw), + c_acc, + static_position=[e], + dynamic_position=[], + ) + + # d_col base for this DC chunk (per-lane via tr_col_half/tr_col_sub) + d_col_base = ( + fx.Index(dc * DC_CHUNK) + + tr_col_half * fx.Index(16) + + tr_col_sub * fx.Index(4) + ) + + for pks in range_constexpr(PV_K_STEPS): + # k_row base for this pks step (per-lane via lane_div_32/tr_k_group) + k_row_base = ( + fx.Index(pks * PV_K_STEP) + + lane_div_32 * fx.Index(4) + + tr_k_group + ) + + # V A-frag via ds_read_tr16_b64: two reads combine into v8f16 + v_lds_lo_f16 = ( + v_lds_base_bytes // fx.Index(2) + + k_row_base * fx.Index(_V_STRIDE) + + d_col_base + ) + v_lds_lo_byte = v_lds_lo_f16 * fx.Index(2) + v_lds_hi_byte = v_lds_lo_byte + fx.Index( + 8 * _V_STRIDE * 2 + ) # +8 toks + + lo_ptr = buffer_ops.create_llvm_ptr( + fx.Int64(v_lds_lo_byte), address_space=3 + ) + hi_ptr = buffer_ops.create_llvm_ptr( + fx.Int64(v_lds_hi_byte), address_space=3 + ) + v_lo_v4 = rocdl.ds_read_tr16_b64( + v4f16_type, lo_ptr + ).result # k=0..3 + v_hi_v4 = rocdl.ds_read_tr16_b64( + v4f16_type, hi_ptr + ).result # k=4..7 + # Combine into v8f16 A-frag: [lo[0..3], hi[0..3]] + v_frag = vector.shuffle(v_lo_v4, v_hi_v4, [0, 1, 2, 3, 4, 5, 6, 7]) + + # P B-frag must match V A-frag tok order (hi-read covers toks +8..11). + p_frag = zero_v8h + for j in range_constexpr(8): + tok_j = ( + fx.Int32(pks * PV_K_STEP) + + lane_div_32 * fx.Int32(4) + + fx.Int32(j % 4) + + fx.Int32((j // 4) * 8) + ) + p_slot = warp_id * fx.Int32(TILE_N) + tok_j + pf = fx.Vector.load(T.vec(1, T.f32), p_lds, [fx.Index(p_slot)])[ + 0 + ] + p_f16 = arith.truncf( + _FX_KV.ir_type, arith.unwrap(fx.Float32(pf)) + ) + p_frag = vector.insert( + p_f16, p_frag, static_position=[j], dynamic_position=[] + ) + + # PV MFMA: A=V_T (tr16), B=P (broadcast) -> C[m=d_sub, n=*]=PV[d] + c_acc = _mfma_pv(T.vec(16, T.f32), [v_frag, p_frag, c_acc, 0, 0, 0]) + + for e in range_constexpr(16): + new_pv_scalars.append( + vector.extract(c_acc, static_position=[e], dynamic_position=[]) + ) + + pv_scalars = new_pv_scalars + state_out = [arith.unwrap(new_max), arith.unwrap(new_sum)] + list( + pv_scalars + ) + results = yield state_out + + final_max = fx.Float32(results[0]) + final_sum = fx.Float32(results[1]) + final_pv_sc = [results[2 + i] for i in range(_N_PV)] + + safe_sum = fx.Float32( + arith.select( + arith.unwrap(final_sum) > c_zero, arith.unwrap(final_sum), c_one + ) + ) + inv_sum = rcp_f32(safe_sum) + out_base = b_idx * stride_qb + g_idx * stride_qg + hq_abs * stride_qh + + if const_expr(_SPLIT): + _pm_base = ( + b_idx * (num_g * split_total * num_hq) + + g_idx * (split_total * num_hq) + + split_idx * num_hq + + hq_abs + ) + _po_base = _pm_base * fx.Int32(_HEAD) + + # Output: C[e] at lane l -> d = dc*32 + ld32*4 + (e//4)*8 + (e%4). + if hq_abs < num_hq: + inv_raw = arith.unwrap(inv_sum) + for dc in range_constexpr(_D_CHUNKS): + for e in range_constexpr(16): + d_out = ( + fx.Int32(dc * DC_CHUNK) + + lane_div_32 * fx.Int32(4) + + fx.Int32((e // 4) * 8 + (e % 4)) + ) + pv_val = final_pv_sc[dc * 16 + e] + if const_expr(_SPLIT): + buffer_ops.buffer_store(pv_val, out_rsrc, _po_base + d_out) + else: + out_val = _FX_OUT( + arith.unwrap(fx.Float32(arith.mulf(pv_val, inv_raw))) + ) + buffer_ops.buffer_store( + arith.unwrap(out_val), out_rsrc, out_base + d_out + ) + + if const_expr(_SPLIT): + if lane == fx.Int32(0): + if hq_abs < num_hq: + buffer_ops.buffer_store(arith.unwrap(final_max), pm_rsrc, _pm_base) + buffer_ops.buffer_store(arith.unwrap(final_sum), ps_rsrc, _pm_base) + + return pa_decode_gfx950_coop_kernel, alloc + + +@functools.lru_cache(maxsize=256) +def _make_gfx950_coop_jit_launcher(head_size, kv_dtype_str, out_dtype_str, split_k): + kernel, _alloc = compile_pa_decode_gfx950_coop( + head_size=head_size, + kv_dtype_str=kv_dtype_str, + output_dtype_str=out_dtype_str, + split_k=split_k, + ) + + @flyc.jit + def _launcher( + out_ptr, + pm_ptr, + ps_ptr, + q_ptr, + k_ptr, + v_ptr, + seq_ptr, + stride_qb, + stride_qg, + stride_qh, + stride_kb, + stride_km, + stride_kg, + stride_kh, + num_hq, + num_g, + kv_max, + num_hkv, + scale, + split_total, + grid_x, + stream: fx.Stream = fx.Stream(None), + ): + from flydsl._mlir import ir as _ir + from flydsl.compiler.kernel_function import CompilationContext + + _alloc.finalized = False + ctx = CompilationContext.get_current() + with _ir.InsertionPoint(ctx.gpu_module_body): + _alloc.finalize() + kernel( + out_ptr, + pm_ptr, + ps_ptr, + q_ptr, + k_ptr, + v_ptr, + seq_ptr, + stride_qb, + stride_qg, + stride_qh, + stride_kb, + stride_km, + stride_kg, + stride_kh, + num_hq, + num_g, + kv_max, + num_hkv, + scale, + split_total, + ).launch(grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), stream=stream) + + return _launcher + + +def pa_decode_gfx950_coop_launch( + Q, K, V, seq_positions, softmax_scale, split_k=0, output_dtype=None +): + """ds_read_tr16_b64 HW transpose for V reads — 8× fewer LDS instructions than scalar reads.""" + from flydsl.runtime.device import get_rocm_arch + from mslk.flydsl.jit import run_compiled + + from .pa_decode_dense import auto_split_k_coop + + B, _, G, H_q, D = Q.shape + _, KV_MAX, _, H_kv, _ = K.shape + # Requires gfx950 and coop-DMA coherence: all NUM_WARPS warps share one K/V LDS + # tile, so GQA ratio must be a multiple of NUM_WARPS. Else fall back to generic. + _coop_ok = ( + H_q % H_kv == 0 and (H_q // H_kv) % NUM_WARPS == 0 and H_q % NUM_WARPS == 0 + ) + if not _coop_ok or not get_rocm_arch().startswith("gfx950"): + from .pa_decode_generic import pa_decode_generic_launch + + return pa_decode_generic_launch( + Q, K, V, seq_positions, softmax_scale, split_k, output_dtype + ) + assert D % MFMA_K_QK == 0 and D % DC_CHUNK == 0 + assert K.dtype in (torch.float16, torch.bfloat16) + if output_dtype is None: + output_dtype = Q.dtype + kv_str = {torch.float16: "f16", torch.bfloat16: "bf16"}[K.dtype] + out_str = {torch.float16: "f16", torch.bfloat16: "bf16", torch.float32: "f32"}[ + output_dtype + ] + if seq_positions is None: + seq_positions = torch.full((B,), KV_MAX, dtype=torch.int32, device=Q.device) + elif seq_positions.dtype != torch.int32: + seq_positions = seq_positions.to(torch.int32) + if split_k == 0: + split_k = auto_split_k_coop(B, G, H_q, KV_MAX) + hq_blocks = (H_q + NUM_WARPS - 1) // NUM_WARPS + out = torch.empty((B, 1, G, H_q, D), dtype=output_dtype, device=Q.device) + sq = Q.stride() + sk2 = K.stride() + dev = Q.device + # Thread the live stream into .launch so the kernel is captured under CUDA graphs + # (a default-stream launch would capture empty). + stream = torch.cuda.current_stream() + if split_k == 1: + dummy = torch.empty(0, dtype=torch.float32, device=dev) + launcher = _make_gfx950_coop_jit_launcher(D, kv_str, out_str, 1) + run_compiled( + launcher, + out, + dummy, + dummy, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + softmax_scale, + split_k, + B * G * hq_blocks, + stream, + ) + else: + po = torch.empty((B, G, split_k, H_q, D), dtype=torch.float32, device=dev) + pm = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + ps = torch.empty((B, G, split_k, H_q), dtype=torch.float32, device=dev) + launcher = _make_gfx950_coop_jit_launcher(D, kv_str, "f32", split_k) + run_compiled( + launcher, + po, + pm, + ps, + Q, + K, + V, + seq_positions, + sq[0], + sq[2], + sq[3], + sk2[0], + sk2[1], + sk2[2], + sk2[3], + H_q, + G, + KV_MAX, + H_kv, + softmax_scale, + split_k, + B * G * hq_blocks * split_k, + stream, + ) + pa_decode_reduce(po, pm, ps, out.squeeze(1), stream=stream) + return out diff --git a/mslk/attention/fmha/flydsl/pa_decode_reduce.py b/mslk/attention/fmha/flydsl/pa_decode_reduce.py new file mode 100644 index 00000000..b8ef3c47 --- /dev/null +++ b/mslk/attention/fmha/flydsl/pa_decode_reduce.py @@ -0,0 +1,385 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""FlyDSL split-K combine (reduce) kernel for paged-attention decode. + +Inputs (from decode kernel): + partial_out : [B, G, max_parts, H_q, D] f32 (UN-normalized numerator sum(p*v)) + partial_max : [B, G, max_parts, H_q] f32 (per-partition global_max) + partial_sum : [B, G, max_parts, H_q] f32 (per-partition exp sum) + out : [B, G, H_q, D] target dtype +Grid (B,G,H_q); Block (WARP_SIZE=64,1,1). Lane handles _CHUNKS=D//64 head-dim pos. + +GOTCHA: partial_out is the un-normalized numerator, so combine each partition by +weight w only (NOT w*partial_sum) — the sum is already folded in. + +Fast path (max_parts ≤ 64): lane l owns partition l; warp reduce for global + max/sum; ds_bpermute broadcasts each partition's normalized weight. +Slow path (max_parts > 64): LDS-staged stats, each lane accumulates independently. +""" + +from __future__ import annotations + +import functools +from typing import Any, Dict, List, Tuple + +import flydsl.compiler as flyc # pyre-ignore[21] +import flydsl.expr as fx # pyre-ignore[21] +import torch +from flydsl.expr import ( # pyre-ignore[21] + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import Int32, T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr # pyre-ignore[21] + +from .utils import exp_f32, rcp_f32, WARP_SIZE, wave_reduce_max_f32, wave_reduce_sum_f32 + +_DTYPE_MAP = { + "f32": (torch.float32, fx.Float32), + "f16": (torch.float16, fx.Float16), + "bf16": (torch.bfloat16, fx.BFloat16), +} + + +def _fx_dtype(dtype_str: str): # pyre-ignore[3] + return _DTYPE_MAP[dtype_str][1] + + +# ── Compiled reduce kernel ──────────────────────────────────────────────────── + + +@functools.lru_cache(maxsize=256) +def _compile_reduce( + head_size: int, + max_parts: int, + output_dtype_str: str, + arch: str, +) -> Tuple[Any, Any]: # pyre-ignore[3] + _HEAD = head_size + _MAX_PARTS = max_parts + _FAST = _MAX_PARTS <= WARP_SIZE + _OUT_FX = _fx_dtype(output_dtype_str) + _CHUNKS = _HEAD // WARP_SIZE + + allocator = SmemAllocator( + None, + arch=arch, + global_sym_name=f"pa_red_p{_MAX_PARTS}_h{_HEAD}_{output_dtype_str}", + ) + if not _FAST: + allocator.ptr = 2 * _MAX_PARTS * 4 # max + sum, f32 each + + @flyc.kernel(known_block_size=(WARP_SIZE, 1, 1)) + def _kernel( + output_ptr: fx.Tensor, + partial_out_ptr: fx.Tensor, + partial_max_ptr: fx.Tensor, + partial_sum_ptr: fx.Tensor, + # partial_out strides: [B, G, SK, Hq, D] + s_po_b: Int32, + s_po_g: Int32, + s_po_part: Int32, + s_po_hq: Int32, + # partial_max/sum strides: [B, G, SK, Hq] — Hq innermost (stride=1) + s_pm_b: Int32, + s_pm_g: Int32, + s_pm_part: Int32, + # output strides: [B, G, Hq, D] + s_o_b: Int32, + s_o_g: Int32, + s_o_hq: Int32, + ) -> None: + lane = gpu.thread_idx.x # 0..WARP_SIZE-1 + bid_b = gpu.block_idx.x + bid_g = gpu.block_idx.y + bid_hq = gpu.block_idx.z + + c_zero = arith.constant(0.0, type=T.f32) + c_one = arith.constant(1.0, type=T.f32) + c_neginf = arith.constant(float("-inf"), type=T.f32) + + po_rsrc = buffer_ops.create_buffer_resource(partial_out_ptr, max_size=True) + pm_rsrc = buffer_ops.create_buffer_resource(partial_max_ptr, max_size=True) + ps_rsrc = buffer_ops.create_buffer_resource(partial_sum_ptr, max_size=True) + out_rsrc = buffer_ops.create_buffer_resource(output_ptr, max_size=True) + + # hq has stride 1 in [B,G,SK,Hq] + pm_base = bid_b * s_pm_b + bid_g * s_pm_g + bid_hq + po_base_hq = bid_b * s_po_b + bid_g * s_po_g + bid_hq * s_po_hq + o_base = bid_b * s_o_b + bid_g * s_o_g + bid_hq * s_o_hq + + if const_expr(_FAST): + # Lane l owns partition l's statistics. + c_mp = arith.constant(_MAX_PARTS, type=T.i32) + active = lane < c_mp + + pm_off = pm_base + lane * s_pm_part + p_max_r = buffer_ops.buffer_load(pm_rsrc, pm_off, vec_width=1, dtype=T.f32) + p_sum_r = buffer_ops.buffer_load(ps_rsrc, pm_off, vec_width=1, dtype=T.f32) + part_max = arith.select(active, p_max_r, c_neginf) + part_sum = arith.select(active, p_sum_r, c_zero) + + gmax = arith.unwrap(wave_reduce_max_f32(fx.Float32(part_max))) + diff = arith.subf(part_max, gmax) + w_f32 = arith.select(active, arith.unwrap(exp_f32(diff)), c_zero) + gsum = arith.unwrap( + wave_reduce_sum_f32(fx.Float32(arith.mulf(w_f32, part_sum))) + ) + inv_sum = arith.unwrap( + rcp_f32(fx.Float32(arith.select(gsum > c_zero, gsum, c_one))) + ) + + norm_w = arith.mulf(w_f32, inv_sum) + norm_w32 = arith.bitcast(T.i32, norm_w) + + # Lane owns a contiguous _CHUNKS-wide slice, so one vec load replaces + # _CHUNKS scalar loads per partition (coalescing preserved). + base_hd = lane * fx.Int32(_CHUNKS) + accs = [c_zero] * _CHUNKS + for p in range_constexpr(_MAX_PARTS): + src = arith.constant(p * 4, type=T.i32) + wi32 = rocdl.ds_bpermute(T.i32, src, norm_w32) + wf32 = arith.bitcast(T.f32, wi32) + poff = po_base_hq + arith.constant(p, type=T.i32) * s_po_part + vals = buffer_ops.buffer_load( + po_rsrc, poff + base_hd, vec_width=_CHUNKS, dtype=T.f32 + ) + if const_expr(_CHUNKS == 1): + accs[0] = arith.addf(accs[0], arith.mulf(vals, wf32)) + else: + for c in range_constexpr(_CHUNKS): + val = vector.extract( + vals, static_position=[c], dynamic_position=[] + ) + accs[c] = arith.addf(accs[c], arith.mulf(val, wf32)) + + for c in range_constexpr(_CHUNKS): + out_val = _OUT_FX(arith.unwrap(fx.Float32(accs[c]))) + buffer_ops.buffer_store( + arith.unwrap(out_val), out_rsrc, o_base + base_hd + fx.Int32(c) + ) + + else: + smem = allocator.get_base() + lm_lds = SmemPtr(smem, 0, T.f32, shape=(_MAX_PARTS,)).get() + ls_lds = SmemPtr(smem, _MAX_PARTS * 4, T.f32, shape=(_MAX_PARTS,)).get() + + for step in range_constexpr((_MAX_PARTS + WARP_SIZE - 1) // WARP_SIZE): + p = step * WARP_SIZE + lane + if const_expr(p < _MAX_PARTS): + pm_off = pm_base + arith.constant(p, type=T.i32) * s_pm_part + lm = buffer_ops.buffer_load( + pm_rsrc, pm_off, vec_width=1, dtype=T.f32 + ) + ls = buffer_ops.buffer_load( + ps_rsrc, pm_off, vec_width=1, dtype=T.f32 + ) + vector.store( + fx.Vector.from_elements([lm], dtype=fx.Float32), + lm_lds, + [fx.Index(arith.constant(p, type=T.i32))], + ) + vector.store( + fx.Vector.from_elements([ls], dtype=fx.Float32), + ls_lds, + [fx.Index(arith.constant(p, type=T.i32))], + ) + gpu.barrier() + + gmax = c_neginf + for p in range_constexpr(_MAX_PARTS): + v = fx.Vector.load( + T.vec(1, T.f32), lm_lds, [fx.Index(arith.constant(p, type=T.i32))] + )[0] + gmax = arith.maximumf(gmax, arith.unwrap(fx.Float32(v))) + + gsum = c_zero + accs = [c_zero] * _CHUNKS + for p in range_constexpr(_MAX_PARTS): + vm = fx.Vector.load( + T.vec(1, T.f32), lm_lds, [fx.Index(arith.constant(p, type=T.i32))] + )[0] + vs = fx.Vector.load( + T.vec(1, T.f32), ls_lds, [fx.Index(arith.constant(p, type=T.i32))] + )[0] + lm_v = arith.unwrap(fx.Float32(vm)) + ls_v = arith.unwrap(fx.Float32(vs)) + w = arith.unwrap(exp_f32(arith.subf(lm_v, gmax))) + gsum = arith.addf(gsum, arith.mulf(w, ls_v)) + poff = po_base_hq + arith.constant(p, type=T.i32) * s_po_part + for c in range_constexpr(_CHUNKS): + hd = lane + fx.Int32(c * WARP_SIZE) + val = buffer_ops.buffer_load( + po_rsrc, poff + hd, vec_width=1, dtype=T.f32 + ) + accs[c] = arith.addf(accs[c], arith.mulf(val, arith.mulf(w, ls_v))) + + safe = arith.select(gsum > c_zero, gsum, c_one) + for c in range_constexpr(_CHUNKS): + hd = lane + fx.Int32(c * WARP_SIZE) + out_val = _OUT_FX(arith.unwrap(fx.Float32(arith.divf(accs[c], safe)))) + buffer_ops.buffer_store(arith.unwrap(out_val), out_rsrc, o_base + hd) + + return _kernel, allocator + + +def compile_pa_decode_reduce( + *, + head_size: int, + max_parts: int, + output_dtype_str: str = "f32", + arch: str = "", +) -> Any: # pyre-ignore[3] + if not arch: + arch = get_rocm_arch() + kernel, _ = _compile_reduce(head_size, max_parts, output_dtype_str, arch) + return kernel + + +# ── JIT launcher ───────────────────────────────────────────────────────────── + + +@functools.lru_cache(maxsize=256) +def _make_reduce_jit_launcher( + head_size: int, + max_parts: int, + output_dtype_str: str, + arch: str, +): # pyre-ignore[3] + kernel, _alloc = _compile_reduce(head_size, max_parts, output_dtype_str, arch) + _fast = max_parts <= WARP_SIZE + + @flyc.jit + def _launcher( + output_ptr: fx.Tensor, + partial_out_ptr: fx.Tensor, + partial_max_ptr: fx.Tensor, + partial_sum_ptr: fx.Tensor, + s_po_b: Int32, + s_po_g: Int32, + s_po_part: Int32, + s_po_hq: Int32, + s_pm_b: Int32, + s_pm_g: Int32, + s_pm_part: Int32, + s_o_b: Int32, + s_o_g: Int32, + s_o_hq: Int32, + grid_b: Int32, + grid_g: Int32, + grid_hq: Int32, + stream: fx.Stream = fx.Stream(None), + ) -> None: + from flydsl._mlir import ir as _ir # pyre-ignore[21] + from flydsl.compiler.kernel_function import ( # pyre-ignore[21] + CompilationContext, + ) + + if not _fast: + _alloc.finalized = False + ctx = CompilationContext.get_current() + with _ir.InsertionPoint(ctx.gpu_module_body): + _alloc.finalize() + + kernel( + output_ptr, + partial_out_ptr, + partial_max_ptr, + partial_sum_ptr, + s_po_b, + s_po_g, + s_po_part, + s_po_hq, + s_pm_b, + s_pm_g, + s_pm_part, + s_o_b, + s_o_g, + s_o_hq, + ).launch(grid=(grid_b, grid_g, grid_hq), block=(WARP_SIZE, 1, 1), stream=stream) + + return _launcher + + +# ── Host API ───────────────────────────────────────────────────────────────── + + +def pa_decode_reduce( + partial_out: torch.Tensor, # [B, G, max_parts, H_q, D] f32 + partial_max: torch.Tensor, # [B, G, max_parts, H_q] f32 + partial_sum: torch.Tensor, # [B, G, max_parts, H_q] f32 + output: torch.Tensor, # [B, G, H_q, D] target dtype + stream: object = None, +) -> None: + """Combine split-K partitions into the final output (in-place). + + Pass the caller's stream so the reduce is captured on the same stream as the + compute kernel under CUDA graphs; defaults to the current stream. + """ + from mslk.flydsl.jit import run_compiled # pyre-ignore[21] + + B, G, max_parts, H_q, D = partial_out.shape + dtype_str = {torch.float32: "f32", torch.float16: "f16", torch.bfloat16: "bf16"}[ + output.dtype + ] + arch = get_rocm_arch() + launcher = _make_reduce_jit_launcher(D, max_parts, dtype_str, arch) + + if stream is None: + stream = torch.cuda.current_stream() + + po, pm, o = partial_out, partial_max, output + run_compiled( + launcher, + output, + partial_out, + partial_max, + partial_sum, + po.stride(0), + po.stride(1), + po.stride(2), + po.stride(3), + pm.stride(0), + pm.stride(1), + pm.stride(2), + o.stride(0), + o.stride(1), + o.stride(2), + B, + G, + H_q, + stream, + ) + + +# ── AOT interface ───────────────────────────────────────────────────────────── + +AOT_ARCHS: List[str] = ["gfx942", "gfx950"] + +AOT_CONFIGS: List[Dict[str, Any]] = [ + {"head_size": hs, "max_parts": mp, "output_dtype_str": dt} + for hs in (64, 128, 256) + for mp in (1, 2, 4, 8, 16, 32, 64) + for dt in ("f32", "f16", "bf16") +] + + +def compile_aot_config(config: Dict[str, Any], arch: str) -> None: + compile_pa_decode_reduce( + head_size=config["head_size"], + max_parts=config["max_parts"], + output_dtype_str=config["output_dtype_str"], + arch=arch, + ) diff --git a/mslk/attention/fmha/flydsl/utils.py b/mslk/attention/fmha/flydsl/utils.py new file mode 100644 index 00000000..70056f3f --- /dev/null +++ b/mslk/attention/fmha/flydsl/utils.py @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Shared low-level FlyDSL helpers for attention kernels. + +Thin re-export of the arch-generic CDNA primitives in +mslk.flydsl.kernels.common.kernel_intrinsics; kept as a stable import site for the +pa_decode_* kernels. +""" + +from mslk.flydsl.kernels.common.kernel_intrinsics import ( # noqa: F401 + dpp_xor_f32, + exp2_f32, + exp_f32, + extract_global_ptr, + global_load_f16x2, + global_load_f32, + global_load_i64x2, + maxnumf, + mfma_f32_16x16x16_bf16, + mfma_f32_16x16x16_f16, + mfma_f32_16x16x4_f32, + rcp_f32, + select_f32, + smem_bytes, + SMEM_BYTES_GFX942, + SMEM_BYTES_GFX950, + WARP_SIZE, + wave_reduce_max_f32, + wave_reduce_sum_f32, +) +from mslk.flydsl.kernels.common.kernels_common import get_warp_size # noqa: F401 diff --git a/mslk/attention/fmha/flydsl_decoder.py b/mslk/attention/fmha/flydsl_decoder.py new file mode 100644 index 00000000..cd727344 --- /dev/null +++ b/mslk/attention/fmha/flydsl_decoder.py @@ -0,0 +1,166 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# pyre-unsafe + +from typing import Any, Iterable, List, Optional, Set, Tuple + +import torch +from mslk.flydsl.common import require_flydsl + +from .attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask +from .common import AttentionFwOpBase, Context, Inputs +from .flydsl.layout_utils import canonicalize_qkv_5d, normalize_seq_positions +from .utils.op_common import get_operator, register_operator + +# FlyDSL is a mandatory ROCm-only backend but absent on CUDA/CPU builds; guard the +# kernel import so `import mslk.attention.fmha` still works there. require_flydsl() +# in apply() raises a clear error before pa_decode_launch is ever called. +try: + from .flydsl.pa_decode_dense import pa_decode_launch +except ImportError: + pa_decode_launch = None + + +def _flydsl_decode_forward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + seq_positions: Optional[torch.Tensor], + scale: float, +) -> torch.Tensor: + q5, k5, v5 = canonicalize_qkv_5d(query, key, value) + B = q5.shape[0] + KV_MAX = k5.shape[1] + seq = normalize_seq_positions(seq_positions, B, KV_MAX, q5.device) + + # split_k=0 -> kernel's auto split-K heuristic (fills the GPU to hide memory latency). + return pa_decode_launch(q5, k5, v5, seq, scale, split_k=0) + + +@register_operator +class FwOp(AttentionFwOpBase): + """FlyDSL dense decode op (gfx942/gfx950). f16/bf16/f32 KV, dense padded layout, GQA/MQA.""" + + OPERATOR = get_operator("xformers", "efficient_attention_forward_decoder_ck") + SUPPORTED_DEVICES: Set[str] = {"cuda"} + SUPPORTED_DTYPES: Set[torch.dtype] = {torch.half, torch.bfloat16} + # pyrefly: ignore [bad-override-mutable-attribute] + SUPPORTED_MAX_K: int = 256 + SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = ( + type(None), + BlockDiagonalCausalWithOffsetPaddedKeysMask, + ) + SUPPORTS_DROPOUT = False + SUPPORTS_CUSTOM_SCALE = True + SUPPORTS_BMGHK = True + NAME = "flydsl_decoderF" + + @classmethod + def not_supported_reasons(cls, d: Inputs) -> List[str]: # noqa: C901 + reasons = super(FwOp, cls).not_supported_reasons(d) + + attn_bias = d.attn_bias + if isinstance(attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask): + if d.query.shape[0] != 1: + reasons.append( + f"One formal batch element expected; got {d.query.shape[0]}" + ) + + if d.query.shape[-1] > cls.SUPPORTED_MAX_K: + reasons.append( + f"Got head_dim={d.query.shape[-1]}; only head_dim<={cls.SUPPORTED_MAX_K} is supported for now." + ) + + threads_per_warp = 64 # TODO: ideally query the platform here + required_alignment = 0 + head_dim = d.query.shape[-1] + for vec_size in (4, 2, 1): + if head_dim <= vec_size * threads_per_warp: + required_alignment = vec_size + + if not required_alignment: + reasons.append(f"Got head_dim={head_dim} which is too large") + + if head_dim % required_alignment != 0: + reasons.append( + f"Got head_dim={head_dim}; it needs to be divisible by {required_alignment}" + ) + + if d.key.stride(-1) != 1: + reasons.append("expect keys to have last dim contiguous") + + if d.value.stride(-1) != 1: + reasons.append("expect values to have last dim contiguous") + + q_starts = attn_bias.q_seqinfo.seqstart_py + padding = attn_bias.k_seqinfo.padding + bsz = d.key.shape[1] // padding + num_queries = d.query.shape[1] // bsz + + if q_starts != list(range(0, 1 + bsz, num_queries)): + reasons.append("expect to have same num_queries in each batch") + if bsz != len(q_starts) - 1: + reasons.append("empty lanes not supported yet") + + if attn_bias.k_seqinfo.padding > 8192: + reasons.append("key padding exceeds 8192") + + return reasons + + @classmethod + def apply( + cls, inp: Inputs, needs_gradient: bool + ) -> Tuple[torch.Tensor, Optional[Context]]: + if needs_gradient: + raise NotImplementedError("backward pass is not supported") + attn_bias = inp.attn_bias + q, k, v = inp.get_qkv_in_bmghk() + if attn_bias is not None: + assert isinstance(attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask) + attn_bias.k_seqinfo.to(k.device) + attn_bias.q_seqinfo.to(q.device) + padding = attn_bias.k_seqinfo.padding + seq_positions_gpu = attn_bias.k_seqinfo.seqlen + else: + padding = k.shape[1] + seq_positions_gpu = None + + if attn_bias is not None: + # key: (1, B * padding, G, 1 if multiquery else Hkv, D) + # value: like key + # query: (1, B * q_seqlen, G, Hq, D) + multiquery = k.stride(3) == 0 + if multiquery: + key = k[0, :, :, :1].unflatten(0, (-1, padding)) + value = v[0, :, :, :1].unflatten(0, (-1, padding)) + else: + key = k[0].unflatten(0, (-1, padding)) + value = v[0].unflatten(0, (-1, padding)) + query = q[0].unflatten(0, (key.shape[0], -1)) + else: + # key: (B, padding, G, 1 if multiquery else Hkv, D) + # value: like key + # query: (B, q_seqlen, G, Hq, D) + key = k + query = q + value = v + + if inp.scale is not None: + qk_scale = inp.scale + else: + qk_scale = torch.rsqrt( + torch.tensor(key.shape[-1], dtype=torch.float32) + ).item() + + require_flydsl() + out = _flydsl_decode_forward( + query=query, + key=key, + value=value, + seq_positions=seq_positions_gpu, + scale=qk_scale, + ) + return out, None diff --git a/mslk/attention/fmha/flydsl_splitk.py b/mslk/attention/fmha/flydsl_splitk.py new file mode 100644 index 00000000..7db3db38 --- /dev/null +++ b/mslk/attention/fmha/flydsl_splitk.py @@ -0,0 +1,222 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# pyre-unsafe + +from typing import Any, Iterable, List, Optional, Tuple + +import torch +from mslk.flydsl.common import is_flydsl_available, require_flydsl + +from .attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask +from .common import AttentionFwOpBase, check_lastdim_alignment_stride1, Context, Inputs +from .flydsl.layout_utils import canonicalize_qkv_5d, normalize_seq_positions +from .utils.op_common import get_operator, register_operator + +# FlyDSL is a mandatory ROCm-only backend but absent on CUDA/CPU builds; guard the +# kernel import so `import mslk.attention.fmha` still works there. require_flydsl() +# in apply() raises a clear error before pa_decode_launch is ever called. +try: + from .flydsl.pa_decode_dense import pa_decode_launch +except ImportError: + pa_decode_launch = None + + +def _flydsl_splitk_forward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + seq_positions: Optional[torch.Tensor], + scale: float, + split_k: int, +) -> torch.Tensor: + q5, k5, v5 = canonicalize_qkv_5d(query, key, value) + B = q5.shape[0] + KV_MAX = k5.shape[1] + seq = normalize_seq_positions(seq_positions, B, KV_MAX, q5.device) + + return pa_decode_launch(q5, k5, v5, seq, scale, split_k=split_k) + + +@register_operator +class FwOp(AttentionFwOpBase): + OPERATOR = get_operator("xformers", "efficient_attention_forward_decoder_splitk_ck") + SUPPORTED_DEVICES = {"cuda"} + SUPPORTED_DTYPES = { + torch.half, + torch.bfloat16, + } # Those are dtypes of Q. In the quantized case K/V has dtype int32 + SUPPORTED_MAX_K = 256 + SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = ( + type(None), + BlockDiagonalCausalWithOffsetPaddedKeysMask, + ) + SUPPORTS_DROPOUT = False + SUPPORTS_CUSTOM_SCALE = True + SUPPORTS_BMGHK = True + NAME = "flydsl_splitKF" + + SPLIT_K: Optional[int] = None + BLOCK_M = 16 + BLOCK_N = 64 + + NUM_GROUPS = 1 # Default quantization is row-wise + + @classmethod + def shape_not_supported_reasons( + cls, Mq: int, Mkv: int, K: int, Kv: int + ) -> List[str]: + reasons = super().shape_not_supported_reasons(Mq, Mkv, K, Kv) + return reasons + + @classmethod + def not_supported_reasons(cls, d: Inputs) -> List[str]: + reasons = super(FwOp, cls).not_supported_reasons(d) + check_lastdim_alignment_stride1(reasons, "query", d.query, 8) + if d.key.dtype != torch.int32: + check_lastdim_alignment_stride1(reasons, "key", d.key, 8) + check_lastdim_alignment_stride1(reasons, "value", d.value, 8) + if not is_flydsl_available(): + reasons.append("FlyDSL is not available for this GPU architecture") + + q_len = d.query.shape[1] + if isinstance(d.attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask): + seqinfo = d.attn_bias.q_seqinfo + if q_len != seqinfo.seqstart_py[-1]: + reasons.append( + f"Expected total {seqinfo.seqstart_py[-1]} queries not {q_len}" + ) + q_len = seqinfo.min_seqlen + if q_len != seqinfo.max_seqlen: + reasons.append( + "Variable query len is not supported in the presence of causal mask." + ) + + if d.key.ndim in [4, 5] and d.key.shape[-2] != 1: + if d.key.stride(-2) == 0 and d.value.stride(-2) == 0 and q_len > 1: + reasons.append("multiquery is only supported with query seqlen=1") + + if d.attn_bias is not None and q_len > 1: + reasons.append( + "query with seqlen > 1 is not supported in the presence of causal mask" + ) + return reasons + + @classmethod + def get_split_k(cls, B: int, H: int, Mk: int) -> int: + """Heuristic for the number of splits""" + bh = max(B * H, 1) # NOTE: Handle B*h=0 case + split_k = max(Mk, 1024) // bh + max_chunk_size = 64 if Mk <= 512 and bh <= 64 else 128 + while split_k > 0 and Mk / split_k < max_chunk_size: + split_k = split_k // 2 + split_k = min(split_k, 64) + split_k = max(split_k, 1) + return split_k + + @classmethod + def apply( + cls, inp: Inputs, needs_gradient: bool + ) -> Tuple[torch.Tensor, Optional[Context]]: + attn_bias = inp.attn_bias + q, k, v = inp.get_qkv_in_bmghk() + + if attn_bias is not None: + assert isinstance(attn_bias, BlockDiagonalCausalWithOffsetPaddedKeysMask) + attn_bias.k_seqinfo.to(k.device) + attn_bias.q_seqinfo.to(q.device) + padding = attn_bias.k_seqinfo.padding + seq_positions_gpu = attn_bias.k_seqinfo.seqlen + else: + padding = k.shape[1] + seq_positions_gpu = None + + if attn_bias is not None: + # key: (1, B * padding, G, 1 if multiquery else Hkv, D) + # value: like key + # query: (1, B * q_seqlen, G, Hq, D) + multiquery = k.stride(3) == 0 + if multiquery: + key = k[0, :, :, :1].unflatten(0, (-1, padding)) + value = v[0, :, :, :1].unflatten(0, (-1, padding)) + else: + key = k[0].unflatten(0, (-1, padding)) + value = v[0].unflatten(0, (-1, padding)) + query = q[0].unflatten(0, (key.shape[0], -1)) + else: + # key: (B, padding, G, 1 if multiquery else Hkv, D) + # value: like key + # query: (B, q_seqlen, G, Hq, D) + key = k + query = q + value = v + + B, _, _, H, _ = query.shape + _, Mk, _, _, _ = key.shape + + if cls.SPLIT_K is not None: + split_k = cls.SPLIT_K + else: + # Use heuristics + split_k = cls.get_split_k(B, H, Mk) + + if inp.scale is not None: + qk_scale = inp.scale + else: + qk_scale = torch.rsqrt( + torch.tensor(k.shape[-1], dtype=torch.float32) + ).item() + + require_flydsl() + out = _flydsl_splitk_forward( + query=query, + key=key, + value=value, + seq_positions=seq_positions_gpu, + scale=qk_scale, + split_k=split_k, + ) + + return out, None + + +class FwOp_S1(FwOp): + SPLIT_K = 1 + NAME = "flydsl_splitK1" + + +class FwOp_S2(FwOp): + SPLIT_K = 2 + NAME = "flydsl_splitK2" + + +class FwOp_S4(FwOp): + SPLIT_K = 4 + NAME = "flydsl_splitK4" + + +class FwOp_S8(FwOp): + SPLIT_K = 8 + NAME = "flydsl_splitK8" + + +class FwOp_S16(FwOp): + SPLIT_K = 16 + NAME = "flydsl_splitK16" + + +class FwOp_S32(FwOp): + SPLIT_K = 32 + NAME = "flydsl_splitK32" + + +class FwOp_S64(FwOp): + SPLIT_K = 64 + NAME = "flydsl_splitK64" + + +class FwOp_S128(FwOp): + SPLIT_K = 128 + NAME = "flydsl_splitK128" diff --git a/mslk/attention/fmha/triton_splitk.py b/mslk/attention/fmha/triton_splitk.py index ea74400c..2d589e07 100644 --- a/mslk/attention/fmha/triton_splitk.py +++ b/mslk/attention/fmha/triton_splitk.py @@ -862,6 +862,14 @@ def grid(META): ) IS_HIP = torch.version.hip is not None + # fp8 byte format must match how the KV cache was quantized: gfx942 e4m3fnuz, + # gfx950/CUDA e4m3fn. Derive from supports_float8_fnuz (NOT hardcoded per-HIP). + if IS_HIP: + from mslk.utils.device import supports_float8_fnuz + + FP8_FNUZ = supports_float8_fnuz(throw_on_hip_incompatibility=False) + else: + FP8_FNUZ = False if inp.quantize_pv_to_fp8: v = v.view(torch.int8) @@ -937,6 +945,7 @@ def grid(META): HAS_ADDITIVE_BIAS=attn_bias_tensor is not None, NUM_PROGRAMS_DIM2_CONST=split_k, IS_HIP=IS_HIP, + FP8_FNUZ=FP8_FNUZ, QUANTIZE_PV_TO_FP8=inp.quantize_pv_to_fp8, QUANTIZE_QK_TO_FP8=inp.quantize_qk_to_fp8, USE_FP32_SCALES=inp.use_fp32_scales, diff --git a/mslk/flydsl/aot.py b/mslk/flydsl/aot.py index 90936a86..ac7213db 100644 --- a/mslk/flydsl/aot.py +++ b/mslk/flydsl/aot.py @@ -34,7 +34,11 @@ # Module paths of AOT-eligible FlyDSL kernel modules. Each must expose # AOT_CONFIGS, AOT_ARCHS, and compile_aot_config(config, arch). -_AOT_KERNEL_MODULES: List[str] = [] +_AOT_KERNEL_MODULES: List[str] = [ + "mslk.attention.fmha.flydsl.pa_decode_dense", + "mslk.attention.fmha.flydsl.pa_decode_reduce", + "mslk.attention.fmha.flydsl.pa_decode_fp8", +] _DEFAULT_MAX_WORKERS: int = 64 diff --git a/mslk/flydsl/kernels/common/kernel_intrinsics.py b/mslk/flydsl/kernels/common/kernel_intrinsics.py new file mode 100644 index 00000000..88bd97dd --- /dev/null +++ b/mslk/flydsl/kernels/common/kernel_intrinsics.py @@ -0,0 +1,215 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Arch-generic CDNA/FlyDSL kernel-authoring primitives (wave64). + +Scalar/vector math intrinsics, DPP cross-lane helpers, wave reductions, global +loads, and MFMA selection wrappers shared across FlyDSL kernels. Not attention +specific — see mslk.attention.fmha.flydsl.utils, which re-exports these. +""" + +import flydsl.expr as fx # pyre-ignore[21] +from flydsl._mlir import ir # pyre-ignore[21] +from flydsl._mlir.dialects import ( # pyre-ignore[21] # pyre-ignore[21] + llvm, + math as mlir_math, +) +from flydsl.expr import arith, buffer_ops, rocdl # pyre-ignore[21] +from flydsl.expr.typing import T # pyre-ignore[21] +from flydsl.runtime.device import get_rocm_arch # pyre-ignore[21] +from flydsl.utils.smem_allocator import SMEM_CAPACITY_MAP # pyre-ignore[21] + + +WARP_SIZE: int = 64 # CDNA wave64 (gfx942, gfx950) + + +def smem_bytes(arch=None) -> int: # pyre-ignore[2] + """LDS capacity in bytes for the given arch (from FlyDSL's known map).""" + if arch is None: + arch = get_rocm_arch() + cap = SMEM_CAPACITY_MAP.get(arch) + if cap is None: + raise ValueError(f"Unsupported arch {arch!r}") + return cap + + +# gfx942 (CDNA3/MI300): 64 KB LDS. gfx950 (CDNA4/MI355): 160 KB LDS. +SMEM_BYTES_GFX942 = 65536 +SMEM_BYTES_GFX950 = 163840 + + +# ── Scalar / vector math intrinsics ───────────────────────────────────────── + + +def rcp_f32(value): # pyre-ignore[2,3] + """Reciprocal via `llvm.amdgcn.rcp.f32` (single instruction).""" + return rocdl.rcp(T.f32, value) + + +def exp_f32(value): # pyre-ignore[2,3] + """Scalar `e^value` via mlir math.exp. Use (not exp2) to match CK natural-exp softmax.""" + raw = ( + arith.unwrap(value) + if hasattr(value, "ir_value") or hasattr(value, "type") + else value + ) + return mlir_math.exp(raw) + + +def exp2_f32(value): # pyre-ignore[2,3] + """Scalar `2^value` via `llvm.amdgcn.exp2.f32` (single v_exp_f32). Used by the + exp2-domain softmax in the MFMA decode kernels.""" + raw = arith.unwrap(value) if hasattr(value, "ir_value") else value + return fx.Float32( + llvm.call_intrinsic(ir.F32Type.get(), "llvm.amdgcn.exp2.f32", [raw], [], []) + ) + + +def maxnumf(a, b): # pyre-ignore[2,3] + """Non-NaN-propagating max — single `v_max_f32` instruction.""" + return type(a)(arith.maxnumf(arith.unwrap(a), arith.unwrap(b))) + + +def select_f32(cond, a, b): # pyre-ignore[2,3] + return arith.select(cond, arith.unwrap(a), arith.unwrap(b)) + + +# ── DPP cross-lane helpers (wave64 CDNA only) ──────────────────────────────── + + +def _dpp_xor_i32_raw(src_i32, offset: int): # pyre-ignore[2,3] + """Butterfly-XOR within a 16-lane DPP row (wave64), offsets 1,2,4,8 only. + + For offsets 16,32 use shuffle_xor/ds_swizzle. DPP control values from AMD ISA. + """ + from flydsl._mlir.dialects import llvm as _llvm # pyre-ignore[21] + from flydsl._mlir.ir import IntegerType # pyre-ignore[21] + + def _upd(src, old, ctrl, rmask, bmask): # pyre-ignore[2,3] + i1_ty = IntegerType.get_signless(1) + bound_false = arith.constant(0, type=i1_ty) + return _llvm.call_intrinsic( + T.i32, + "llvm.amdgcn.update.dpp.i32", + [ + old, + src, + arith.unwrap(arith.constant(ctrl, type=T.i32)), + arith.unwrap(arith.constant(rmask, type=T.i32)), + arith.unwrap(arith.constant(bmask, type=T.i32)), + bound_false, + ], + [], + [], + ) + + if offset == 8: + out = _upd(src_i32, src_i32, 280, 0xF, 0xC) + out = _upd(src_i32, out, 264, 0xF, 0x3) + elif offset == 4: + out = _upd(src_i32, src_i32, 276, 0xF, 0xA) + out = _upd(src_i32, out, 260, 0xF, 0x5) + elif offset == 2: + out = _upd(src_i32, src_i32, 78, 0xF, 0xF) + elif offset == 1: + out = _upd(src_i32, src_i32, 177, 0xF, 0xF) + else: + raise ValueError(f"dpp_xor only supports offsets 1,2,4,8; got {offset}") + return out + + +def dpp_xor_f32(src, offset: int): # pyre-ignore[2,3] + """F32 butterfly-XOR within a 16-lane DPP row (wave64, offsets 1/2/4/8).""" + from flydsl._mlir.dialects import arith as _arith_dialect # pyre-ignore[21] + + raw = arith.unwrap(src) if hasattr(src, "ir_value") else src + src_i32 = _arith_dialect.BitcastOp(T.i32, raw).result + out_i32 = _dpp_xor_i32_raw(src_i32, offset) + return fx.Float32(_arith_dialect.BitcastOp(T.f32, out_i32).result) + + +def wave_reduce_max_f32(val): # pyre-ignore[2,3] + """Full wave64 max reduction: DPP XOR (8,4,2,1) then shuffle_xor (32,16).""" + for sh in (8, 4, 2, 1): + val = maxnumf(val, dpp_xor_f32(val, sh)) + c_w = arith.constant(WARP_SIZE, type=T.i32) + for sh in (32, 16): + other = val.shuffle_xor(arith.constant(sh, type=T.i32), c_w) + val = maxnumf(val, fx.Float32(other)) + return val + + +def wave_reduce_sum_f32(val): # pyre-ignore[2,3] + """Full wave64 warp-level sum reduction.""" + for sh in (8, 4, 2, 1): + val = fx.Float32( + arith.addf(arith.unwrap(val), arith.unwrap(dpp_xor_f32(val, sh))) + ) + c_w = arith.constant(WARP_SIZE, type=T.i32) + for sh in (32, 16): + other = val.shuffle_xor(arith.constant(sh, type=T.i32), c_w) + val = fx.Float32(arith.addf(arith.unwrap(val), arith.unwrap(fx.Float32(other)))) + return val + + +# ── Global pointer extraction ──────────────────────────────────────────────── + + +def extract_global_ptr(tensor): # pyre-ignore[2,3] + """Extract a raw `!llvm.ptr<1>` from a FlyDSL tensor argument.""" + from flydsl._mlir.dialects import fly as _fly # pyre-ignore[21] + + raw = ( + tensor.ir_value() + if hasattr(tensor, "ir_value") and not isinstance(tensor, ir.Value) + else tensor + ) + ptr_type = ir.Type.parse("!llvm.ptr<1>") + return _fly.extract_aligned_pointer_as_index(ptr_type, raw) + + +def global_load_f32(global_ptr, byte_offset_i64): # pyre-ignore[2,3] + """Load one f32 from a raw global pointer + byte offset.""" + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset_i64), elem_type=T.i8 + ) + return llvm.LoadOp(T.f32, ptr, alignment=4).result + + +def global_load_f16x2(global_ptr, byte_offset_i64): # pyre-ignore[2,3] + """Load a packed pair of f16 values (32-bit aligned).""" + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset_i64), elem_type=T.i8 + ) + return llvm.LoadOp(T.i32, ptr, alignment=4).result + + +def global_load_i64x2(global_ptr, byte_offset_i64): # pyre-ignore[2,3] + """Load 128 bits (two i64) from a raw global pointer + byte offset.""" + ptr = buffer_ops.get_element_ptr( + global_ptr, byte_offset=fx.Int64(byte_offset_i64), elem_type=T.i8 + ) + return llvm.LoadOp(T.i64x2, ptr, alignment=16).result + + +# ── MFMA selection helpers ─────────────────────────────────────────────────── + + +def mfma_f32_16x16x16_f16(a, b, acc): # pyre-ignore[2,3] + """f16 × f16 → f32 MFMA (16×16×16).""" + return rocdl.mfma_f32_16x16x16f16(T.f32x4, [a, b, acc, 0, 0, 0]) + + +def mfma_f32_16x16x16_bf16(a, b, acc): # pyre-ignore[2,3] + """bf16 × bf16 → f32 MFMA (16×16×16); uses the 1k (accumulator) variant.""" + return rocdl.mfma_f32_16x16x16bf16_1k(T.f32x4, [a, b, acc, 0, 0, 0]) + + +def mfma_f32_16x16x4_f32(a, b, acc): # pyre-ignore[2,3] + """f32 × f32 → f32 MFMA (16×16×4).""" + return rocdl.mfma_f32_16x16x4f32(T.f32x4, [a, b, acc, 0, 0, 0]) diff --git a/mslk/utils/device.py b/mslk/utils/device.py index 764e3a8b..5486f518 100644 --- a/mslk/utils/device.py +++ b/mslk/utils/device.py @@ -74,6 +74,20 @@ def cuda_version_at_least(major_min: int) -> bool: return int(torch.version.cuda.split(".")[0]) >= major_min +def rocm_version_at_least(major_min: int, minor_min: int = 0) -> bool: + """True on a ROCm build whose HIP runtime version is at least ``(major_min, minor_min)``. + + Parses ``torch.version.hip`` (e.g. ``"7.14.60850"``). Returns ``False`` on CUDA or + CPU-only builds. + """ + if torch.version.hip is None: + return False + parts = torch.version.hip.split(".") + major = int(parts[0]) + minor = int(parts[1]) if len(parts) > 1 else 0 + return (major, minor) >= (major_min, minor_min) + + def get_gfx_arch_name() -> str: """Return the ROCm ``gcnArchName`` of the current device (e.g. ``gfx942``). diff --git a/test/attention/fmha/test_mem_eff_attention.py b/test/attention/fmha/test_mem_eff_attention.py index ba6cba36..a9bd8602 100644 --- a/test/attention/fmha/test_mem_eff_attention.py +++ b/test/attention/fmha/test_mem_eff_attention.py @@ -1049,6 +1049,28 @@ def test_decoder_ck( ) +@rocm_only +@pytest.mark.parametrize("kv_heads", [None, 1, 2], ids=_kv_heads_label) +@pytest.mark.parametrize("bsz,n_heads", [(1, 1), (1, 16), (1, 32), (8, 1), (4, 8)]) +@pytest.mark.parametrize("padding", [32, 4096]) +@pytest.mark.parametrize("dtype", ["f16", "bf16"]) +def test_decoder_flydsl( + n_heads: int, + kv_heads: Optional[int], + padding: int, + bsz: int, + dtype: str, +) -> None: + _test_decoder( + fmha.flydsl_decoder.FwOp, + kv_heads=kv_heads, + n_heads=n_heads, # qheads per kv head + padding=padding, + bsz=bsz, + dtype=dtype, + ) + + @sm100_or_better_only @pytest.mark.parametrize("kv_heads", [1, 2, 16], ids=_kv_heads_label) @pytest.mark.parametrize("n_heads", [1, 4, 16]) @@ -1104,6 +1126,130 @@ def test_ck_splitk_decoder( ) +@rocm_only +@pytest.mark.parametrize( + "op", + [ + fmha.flydsl_splitk.FwOp_S1, + fmha.flydsl_splitk.FwOp_S2, + fmha.flydsl_splitk.FwOp_S4, + ], +) +@pytest.mark.parametrize("dtype", ["f16", "bf16"]) +@pytest.mark.parametrize("kv_heads", [None, 1, 2], ids=_kv_heads_label) +@pytest.mark.parametrize("n_heads", [16]) +@pytest.mark.parametrize("d", [128, 256]) +@pytest.mark.parametrize("padding, bsz", [(32, 8), (4096, 1), (32, 1), (4096, 8)]) +def test_flydsl_splitk_decoder( + op, + kv_heads: Optional[int], + n_heads: int, + padding: int, + bsz: int, + dtype: str, + d: int, +) -> None: + _test_decoder( + op, + kv_heads=kv_heads, + n_heads=n_heads, + padding=padding, + bsz=bsz, + dtype=dtype, + d=d, + ) + + +@rocm_only +@pytest.mark.parametrize("dtype", ["f16", "bf16"]) +@pytest.mark.parametrize("n_heads", [1, 16]) +@pytest.mark.parametrize("kv_heads", [1, 2], ids=lambda x: f"kvh{x}") +@pytest.mark.parametrize("padding, bsz", [(512, 2), (2048, 4), (32, 8)]) +@pytest.mark.parametrize("d", [128, 256]) +def test_flydsl_fp8_decoder( + dtype: str, + n_heads: int, + kv_heads: int, + padding: int, + bsz: int, + d: int, +) -> None: + """Correctness of the FlyDSL native-fp8 paged decode kernel. + + Builds a persistent fp8 paged KV cache with precomputed per-token scales via + ``dense_kv_to_fp8_paged``, then runs ``pa_decode_ps_launch`` directly against it — + the real fp8-cache usage, not per-call quantization. Covers MQA (kv_heads=1) and + GQA (kv_heads>1), compared to a full-precision reference within fp8 tolerance. + """ + from mslk.attention.fmha.flydsl.fp8_paged_cache import dense_kv_to_fp8_paged + from mslk.attention.fmha.flydsl.pa_decode_fp8 import pa_decode_ps_launch + from mslk.attention.fmha.flydsl.pa_decode_fp8_dispatch import ( + is_fp8_paged_decode_available, + ) + + if not is_fp8_paged_decode_available(): + pytest.skip("FlyDSL native-fp8 paged decode unavailable (needs gfx950)") + if d % 16 != 0: + pytest.skip("fp8 kernel requires head_dim % 16 == 0") + if kv_heads > 1 and n_heads == 1: + pytest.skip("GQA needs n_heads (query heads per group) > 1") + + dtype_ = {"f16": torch.float16, "bf16": torch.bfloat16}[dtype] + torch.manual_seed(1) + dev = "cuda" + B, G, Hq, D = bsz, kv_heads, n_heads, d + + # Packed BMGHK for the reference/bias: one KV head per group broadcast to Hq. + k_folded = (1, B * padding, G, 1, D) + k_shape = (1, B * padding, G, Hq, D) + kf = torch.randn(k_folded, dtype=dtype_, device=dev) + vf = torch.randn(k_folded, dtype=dtype_, device=dev) + k = kf.expand(k_shape) + v = vf.expand(k_shape) + q = torch.randn((1, B, G, Hq, D), dtype=dtype_, device=dev) + seq = torch.randint(1, padding + 1, (B,), dtype=torch.int32, device=dev) + + # Build the fp8 paged cache ONCE (precomputed scales), then run the kernel over it. + # Per-batch dense KV [B, padding, G, 1, D] is what the cache builder expects. + k_dense = kf[0].unflatten(0, (B, padding)) + v_dense = vf[0].unflatten(0, (B, padding)) + key_cache, value_cache, key_scale, value_scale, block_tables = ( + dense_kv_to_fp8_paged(k_dense, v_dense, block_size=16) + ) + BG = B * G + context_lengths = seq.view(B, 1).expand(B, G).reshape(BG).contiguous() + q_flat = q.reshape(BG, Hq, D).contiguous() + out = torch.zeros(BG, Hq, D, dtype=q.dtype, device=dev) + pa_decode_ps_launch( + out, + q_flat, + key_cache, + value_cache, + context_lengths, + float(D**-0.5), + key_scale=key_scale, + value_scale=value_scale, + block_tables=block_tables, + max_context_partition_num=0, + ) + + attn_bias = fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask.from_seqlens( + q_seqlen=[1] * B, + kv_seqlen=seq.tolist(), + kv_padding=padding, + ) + ref_output = ref_attention_for_test(q, k, v, attn_bias) + out = out.reshape(ref_output.shape) + + # fp8 (e4m3fn) has ~2 mantissa bits -> loose tolerance vs the full-precision ref. + assert_allclose( + out.to(ref_output.dtype), + ref_output, + atol=0.2, + rtol=0.15, + ) + + @sm80_or_better_only @pytest.mark.parametrize( "op", @@ -1997,13 +2143,19 @@ def test_triton_splitk_rowwise_fp8( inp_ref, op=fmha.triton_splitk.FwOp ) + # ROCm gfx950 OCP e4m3fn snaps a few values differently than the fnuz grid, so + # loosen tolerances on ROCm only; CUDA keeps the original tight values. + is_hip = torch.version.hip is not None atol = 5e-3 - if Hkv == 2 and torch.version.hip is not None: - # XXX why is this needed? + rtol = 5e-3 + if Hkv == 2 and is_hip: atol = 1e-2 - torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=5e-3) + torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=rtol) assert context_fp8 is not None and context_ref is not None - torch.testing.assert_close(context_fp8.lse, context_ref.lse, atol=5e-4, rtol=5e-4) + lse_tol = 5e-3 if is_hip else 5e-4 + torch.testing.assert_close( + context_fp8.lse, context_ref.lse, atol=lse_tol, rtol=lse_tol + ) # Paged K/V cache @@ -2017,8 +2169,12 @@ def test_triton_splitk_rowwise_fp8( ) = fmha._memory_efficient_attention_forward_requires_grad( inp_fp8_paged, op=fmha.triton_splitk.FwOp ) + # Non-paged vs paged fp8 output: a few elements snap to a different e4m3 grid point + # between the two layouts on ROCm; CUDA keeps the original tight value. The LSE is + # identical between layouts on both platforms, so it keeps the tight tolerance. + paged_tol = 5e-3 if is_hip else 2e-3 torch.testing.assert_close( - attn_output_fp8, attn_output_fp8_paged, atol=2e-3, rtol=2e-3 + attn_output_fp8, attn_output_fp8_paged, atol=paged_tol, rtol=paged_tol ) assert context_fp8_paged is not None torch.testing.assert_close( diff --git a/test/attention/fmha/utils.py b/test/attention/fmha/utils.py index fc475aa1..9e88ecc7 100644 --- a/test/attention/fmha/utils.py +++ b/test/attention/fmha/utils.py @@ -25,6 +25,7 @@ ref_attention_bmhk, ) from mslk.attention.fmha.triton_splitk import InputsFp8 +from mslk.utils.triton.fp8_utils import get_fp8_constants IN_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER") is not None @@ -186,9 +187,9 @@ def construct_fp8_attention_inputs( k = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device) v = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device) - pt_fp8_dtype = ( - torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn - ) + # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn, + # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN. + pt_fp8_dtype = get_fp8_constants()[0] qfn = quantize_fp8_symmetric if use_symmetric else quantize_fp8_asymmetric @@ -427,9 +428,9 @@ def add_q_fp8_to_inputs( InputsFp8 object with quantized query tensor """ inp.quantize_qk_to_fp8 = True - pt_fp8_dtype = ( - torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn - ) + # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn, + # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN. + pt_fp8_dtype = get_fp8_constants()[0] # Get original query tensor q = inp.query