diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index e5467f3ed30..91a2e18bfb8 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -10,14 +10,14 @@ from packaging.version import Version -from .utils import is_flydsl_available from .moe_common import GateMode +from .utils import is_flydsl_available _MIN_FLYDSL_VERSION = Version("0.2.4") __all__ = [ - "is_flydsl_available", "GateMode", + "is_flydsl_available", ] if is_flydsl_available(): @@ -39,6 +39,7 @@ from .fmha_kernels import flydsl_flash_attn_func from .gemm_kernels import flydsl_hgemm, flydsl_preshuffle_gemm_a8 + from .hstu_attention_kernels import flydsl_hstu_attention_fwd from .kernels.fp8_mqa_logits import ( DEFAULT_VARIANT as FP8_MQA_LOGITS_DEFAULT_VARIANT, ) @@ -68,6 +69,7 @@ "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_hgemm", + "flydsl_hstu_attention_fwd", "flydsl_moe_stage1", "flydsl_moe_stage2", "flydsl_pa_mqa_logits_fp4", diff --git a/aiter/ops/flydsl/hstu_attention_kernels.py b/aiter/ops/flydsl/hstu_attention_kernels.py new file mode 100644 index 00000000000..9760da576fa --- /dev/null +++ b/aiter/ops/flydsl/hstu_attention_kernels.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""High-level FlyDSL HSTU Attention Forward API.""" + +from __future__ import annotations + +import csv +import functools +from collections.abc import Callable +from math import gcd, lcm +from pathlib import Path + +import flydsl.expr as fx +import torch +from flydsl.runtime.device import get_rocm_arch + +from aiter import logger +from aiter.ops.flydsl.kernels.hstu_attention_fwd import ( + MFMA_M, + WARP_SIZE, + _arch_dma_params, + build_hstu_attention_fwd, + validate_hstu_attention_fwd, +) +from aiter.ops.triton.utils.common_utils import prev_power_of_2 + +from .kernels.tensor_shim import _run_compiled, get_dtype_str + +__all__ = [ + "flydsl_hstu_attention_fwd", +] + + +_GPU_ARCH = get_rocm_arch() + + +# Tuned kernel configs +# list of column names in the tuned csv file +_CSV_COLUMNS: list[str] = [ + "arch", + "dtype", + "num_heads", + "head_dim", + "hidden_dim", + "batch", + "max_seq_len", + "has_window", + "has_contextual", + "has_targets", + "block_m", + "block_n", + "num_waves", + "waves_per_eu", + "duration", +] + + +def _str2bool(v: bool | str) -> bool: + """Local, dependency-free bool coercion for CSV/string flags. + + Intentionally not imported from ``aiter.utility.dtypes`` to keep this module + importable during the setup.py AOT phase: that module pulls in + ``aiter.ops.enum``, which instantiates compiled bindings not yet built at + setup time. + """ + if isinstance(v, bool): + return v + if v.strip().lower() in ("yes", "true", "t", "y", "1"): + return True + if v.strip().lower() in ("no", "false", "f", "n", "0"): + return False + raise ValueError(f"Boolean value expected, got {v!r}") + + +def _problem_key( + arch: str, + dtype: str, + num_heads: int, + head_dim: int, + hidden_dim: int, + batch: int, + max_seq_len: int, + has_window: bool | str, + has_contextual: bool | str, + has_targets: bool | str, +) -> tuple: + return ( + arch.strip().lower(), + dtype.strip().lower(), + int(num_heads), + int(head_dim), + int(hidden_dim), + prev_power_of_2(int(batch)), + prev_power_of_2(int(max_seq_len)), + _str2bool(has_window), + _str2bool(has_contextual), + _str2bool(has_targets), + ) + + +@functools.lru_cache +def _tuned_config_map(tuned_file: str | None = None) -> dict[tuple, dict]: + def _parse_row(row: dict) -> tuple[tuple, float, dict]: + if set(row.keys()) != set(_CSV_COLUMNS): + raise KeyError(f"unexpected columns: {set(row.keys()) ^ set(_CSV_COLUMNS)}") + + duration = float(row["duration"]) + + problem_key = _problem_key( + row["arch"], + row["dtype"], + row["num_heads"], + row["head_dim"], + row["hidden_dim"], + row["batch"], + row["max_seq_len"], + (row["has_window"]), + row["has_contextual"], + row["has_targets"], + ) + kernel_config = { + "block_m": int(row["block_m"]), + "block_n": int(row["block_n"]), + "num_waves": int(row["num_waves"]), + "waves_per_eu": int(row["waves_per_eu"]), + } + return ( + problem_key, + duration, + kernel_config, + ) + + default_tuned_file = Path(__file__).resolve().parent / "hstu_attention_tuned.csv" + + tuned_file: Path = Path(tuned_file) if tuned_file else default_tuned_file + if not tuned_file.is_file(): + return {} + + config_map: dict = {} + with tuned_file.open(mode="r", encoding="utf-8") as f: + for row_idx, row in enumerate(csv.DictReader(f)): + try: + problem_key, duration, kernel_config = _parse_row(row) + except (KeyError, ValueError, TypeError) as exc: + logger.warning( + f"[FlyDSL HSTU Fwd] skipping invalid tuned row {row_idx} in {tuned_file}: {exc}" + ) + continue + + if duration <= 0.0: + continue + + if problem_key not in config_map or duration < config_map[problem_key][0]: + config_map[problem_key] = (duration, kernel_config) + + return { + problem_key: kernel_config + for problem_key, (_, kernel_config) in config_map.items() + } + + +def _get_tuned_config( + *, + dtype_str: str, + num_heads: int, + head_dim: int, + hidden_dim: int, + batch: int, + max_seq_len: int, + max_attn_len: int, + contextual_seq_len: int, + has_targets: bool, +) -> dict: + """ + Returns the tuned kernel config if it exists for the given parameters. + """ + + problem_key = _problem_key( + _GPU_ARCH, + dtype_str, + num_heads, + head_dim, + hidden_dim, + batch, + max_seq_len, + max_attn_len > 0, + contextual_seq_len > 0, + has_targets, + ) + + return _tuned_config_map().get(problem_key, {}) + + +def _cdna4_safe_block_n( + *, + block_n: int, + num_waves: int, + head_dim: int, + hidden_dim: int, + arch: str | None, +) -> int: + """Round ``block_n`` up so the K/V LDS tiles divide the CDNA4 DMA pass. + + CDNA4 (gfx950) uses a ``dwordx4`` ``buffer_load_lds`` (16 B/lane), 4x wider than + CDNA3's ``dword``. A MI300X-derived ``block_n`` can therefore leave the K/V LDS + tiles indivisible by the wider DMA pass (see ``validate_hstu_attention_fwd``), + e.g. ``block_n=48`` with ``hidden_dim=64``. This bumps ``block_n`` to the nearest + multiple that keeps both tiles valid. + + This is a strict no-op on non-gfx950 arches, so the MI300X-tuned block sizes are + preserved exactly. + """ + if not (arch or "").startswith("gfx950"): + return block_n + + _, dma_elems, _, _ = _arch_dma_params(arch) + elems_per_dma_pass = num_waves * WARP_SIZE * dma_elems + head_dim_k = ((head_dim + 63) // 64) * 64 + # block_n must satisfy (block_n * tile_width) % elems_per_dma_pass == 0 for both the + # K tile (head_dim_k) and V tile (hidden_dim); solve for the smallest granularity. + req_k = elems_per_dma_pass // gcd(elems_per_dma_pass, head_dim_k) + req_v = elems_per_dma_pass // gcd(elems_per_dma_pass, hidden_dim) + # keep block_n a multiple of MFMA_M as the validator independently requires. + step = lcm(req_k, req_v, MFMA_M) + return ((block_n + step - 1) // step) * step + + +def _get_default_config( + *, + batch: int, + head_dim: int, + hidden_dim: int, + num_heads: int, + max_seq_len: int, + max_attn_len: int, + arch: str | None = None, +) -> dict: + """ + Heuristic config for when tuning is unavailable. + + Derived from a device sweep over shapes on MI300X. + - block_m by occupancy + - block_n by head/hidden dim (bounded by the LDS K+V tile) + - waves_per_eu lifts residency only for tiny-dim tiles that benefit. + + The block sizes below are MI300X (gfx942/CDNA3) values. On CDNA4 (gfx950) the DMA + pass is 4x wider, so ``block_n`` is rounded up by ``_cdna4_safe_block_n`` for that + arch only; the gfx942 path is returned unchanged. + """ + + def as_dict( + block_m: int, + block_n: int, + num_waves: int, + waves_per_eu: int, + /, + ) -> dict: + return { + "block_m": block_m, + "block_n": block_n, + "num_waves": num_waves, + "waves_per_eu": waves_per_eu, + } + + def _mi300x_config() -> dict: + # hidden_dims 96/160/224 don't divide the K/V DMA pass with num_waves=4 + # This map is so the kernel still runs for these values. + non_64_divisible_map = { + 96: (96, 48, 3, 0), + 160: (160, 80, 5, 0), + 192: (96, 48, 3, 0), + 224: (112, 112, 7, 0), + } + if hidden_dim in non_64_divisible_map: + return as_dict(*non_64_divisible_map[hidden_dim]) + + grid = batch * num_heads * ((max_seq_len + 127) // 128) + dim = max(head_dim, hidden_dim) + + if dim <= 64: + if max_attn_len: + return as_dict(128, 32, 4, 0) + + if grid >= 6144: + return as_dict(256, 32, 4, 0) + if grid >= 768: + return as_dict(128, 32, 4, 2) + return as_dict(64, 32, 4, 2) + + # dim > 64 + if grid >= 2560 or max_seq_len >= 16384: + return as_dict(192, 48, 4, 0) + if grid >= 768: + return as_dict(128, 64, 4, 2) + return as_dict(64, 64, 4, 2) + + config = _mi300x_config() + config["block_n"] = _cdna4_safe_block_n( + block_n=config["block_n"], + num_waves=config["num_waves"], + head_dim=head_dim, + hidden_dim=hidden_dim, + arch=arch if arch is not None else _GPU_ARCH, + ) + return config + + +@functools.lru_cache(maxsize=16384) +def _compile_launcher( + *, + batch: int, + max_seq_len: int, + num_heads: int, + head_dim: int, + hidden_dim: int, + causal: bool, + has_targets: bool, + alpha: float, + max_attn_len: int, + contextual_seq_len: int, + dtype_str: str, + block_m: int | None, + block_n: int | None, + num_waves: int | None, + waves_per_eu: int | None, +) -> tuple[str, Callable]: + # Config overrides (if provided) + custom_config: dict = { + "block_m": block_m, + "block_n": block_n, + "num_waves": num_waves, + "waves_per_eu": waves_per_eu, + } + custom_config = {k: v for k, v in custom_config.items() if v is not None} + + # Tuned config entry + tuned_config = _get_tuned_config( + dtype_str=dtype_str, + num_heads=num_heads, + head_dim=head_dim, + hidden_dim=hidden_dim, + batch=batch, + max_seq_len=max_seq_len, + max_attn_len=max_attn_len, + contextual_seq_len=contextual_seq_len, + has_targets=has_targets, + ) + + # Default hueristic config + default_config = _get_default_config( + batch=batch, + head_dim=head_dim, + hidden_dim=hidden_dim, + num_heads=num_heads, + max_seq_len=max_seq_len, + max_attn_len=max_attn_len, + ) + + kernel_config = { + **default_config, + **tuned_config, + **custom_config, + } + + kwargs: dict = dict( + num_heads=num_heads, + head_dim=head_dim, + hidden_dim=hidden_dim, + batch=batch, + causal=causal, + max_attn_len=max_attn_len, + has_targets=has_targets, + alpha=alpha, + dtype_str=dtype_str, + max_seq_len=max_seq_len, + contextual_seq_len=contextual_seq_len, + **kernel_config, + ) + validate_hstu_attention_fwd(**kwargs) + launcher = build_hstu_attention_fwd(**kwargs) + return launcher + + +def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_offsets: torch.Tensor, + num_targets: torch.Tensor | None, +) -> tuple[int, int, int, int, str]: + tensors: dict[str, torch.Tensor] = { + "q": q, + "k": k, + "v": v, + "seq_offsets": seq_offsets, + } + if num_targets is not None: + tensors["num_targets"] = num_targets + + if not all(t.is_cuda for t in tensors.values()): + raise ValueError("flydsl_hstu_attention_fwd requires device tensors") + if not all(t.device == tensors["q"].device for t in tensors.values()): + raise ValueError("tensors must reside on the same device") + + if q.dim() != 3 or k.dim() != 3 or v.dim() != 3: + raise ValueError( + "q/k/v must be rank 3, got " + f"q={tuple(q.shape)} k={tuple(k.shape)} v={tuple(v.shape)}" + ) + if q.shape != k.shape: + raise ValueError( + "q and k must have the same shape, got " + f"q={tuple(q.shape)} k={tuple(k.shape)}" + ) + if v.shape[0] != q.shape[0] or v.shape[1] != q.shape[1]: + raise ValueError( + "v must share q's token count and head count, got " + f"q={tuple(q.shape)} v={tuple(v.shape)}" + ) + if not (q.dtype == k.dtype == v.dtype): + raise ValueError( + f"q/k/v must share the same dtype, got q={q.dtype} k={k.dtype} v={v.dtype}" + ) + + dtype_str = get_dtype_str(q.dtype) + num_heads, head_dim = q.shape[-2:] + hidden_dim = v.shape[2] + batch = seq_offsets.numel() - 1 + + if dtype_str is None: + raise ValueError(f"Unsupported dtype: get_dtype_str({q.dtype}) is None") + if num_targets is not None: + if num_targets.device != q.device: + raise ValueError( + f"num_targets must be on q's device ({q.device}), got {num_targets.device}" + ) + if num_targets.numel() != batch: + raise ValueError( + f"num_targets length ({num_targets.numel()}) must equal batch ({batch})" + ) + + return ( + batch, + num_heads, + head_dim, + hidden_dim, + dtype_str, + ) + + +def flydsl_hstu_attention_fwd( + N: int, + alpha: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + seq_offsets: torch.Tensor, + causal: bool, + num_targets: torch.Tensor | None, + max_attn_len: int, + contextual_seq_len: int, + *, + block_m: int | None = None, + block_n: int | None = None, + num_waves: int | None = None, + waves_per_eu: int | None = None, + stream: torch.cuda.Stream | None = None, +) -> torch.Tensor: + batch, num_heads, head_dim, hidden_dim, dtype_str = _validate_inputs( + q=q, + k=k, + v=v, + seq_offsets=seq_offsets, + num_targets=num_targets, + ) + + launcher = _compile_launcher( + batch=batch, + max_seq_len=N, + num_heads=num_heads, + head_dim=head_dim, + hidden_dim=hidden_dim, + causal=causal, + has_targets=num_targets is not None, + alpha=alpha, + max_attn_len=max_attn_len, + contextual_seq_len=contextual_seq_len, + dtype_str=dtype_str, + block_m=block_m, + block_n=block_n, + num_waves=num_waves, + waves_per_eu=waves_per_eu, + ) + + out = torch.empty_like(v) + if num_targets is None: + num_targets = torch.zeros(1, dtype=seq_offsets.dtype, device=out.device) + + launch_stream = torch.cuda.current_stream(q.device) if stream is None else stream + with torch.cuda.device(q.device.index): + _run_compiled( + launcher, + q.contiguous(), + k.contiguous(), + v.contiguous(), + seq_offsets.contiguous(), + num_targets.contiguous(), + out, + fx.Stream(launch_stream), + ) + return out diff --git a/aiter/ops/flydsl/kernels/hstu_attention_common.py b/aiter/ops/flydsl/kernels/hstu_attention_common.py new file mode 100644 index 00000000000..9814dafb7d9 --- /dev/null +++ b/aiter/ops/flydsl/kernels/hstu_attention_common.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Shared indexing helpers for the FlyDSL HSTU attention kernels. + +The forward and backward kernels use the same address-computation idioms with +operand roles relabelled. These are factored here so there is a single source of +truth, and so the thread-coordinate decomposition is expressed with FlyDSL +layout algebra (idx2crd over a make_layout) rather than hand-rolled integer +division/modulo. + +All helpers build FlyDSL expressions and must be called from inside a +@flyc.kernel body. +""" + +from __future__ import annotations + +import flydsl.expr as fx + + +def decode_lane(tid, num_waves: int, warp_size: int, mfma_n: int): + """Decompose a flat thread id into (wave_id, lane, lane_div_n, lane_mod_n). + + Uses layout algebra: tid indexes a (num_waves, warp_size) layout to split + wave/lane, and the lane indexes a (warp_size/mfma_n, mfma_n) layout to split + the MFMA lane coordinate. Equivalent to tid//warp_size, tid%warp_size, + lane//mfma_n, lane%mfma_n, expressed as coordinate maps. + """ + # idx2crd/get yield index-typed coordinates; cast back to Int32 to match the + # kernels' i32 address arithmetic. + wave_lane = fx.idx2crd(tid, fx.make_layout((num_waves, warp_size), (warp_size, 1))) + wave_id = fx.Int32(fx.get(wave_lane, 0)) + lane = fx.Int32(fx.get(wave_lane, 1)) + + lane_split = fx.idx2crd( + lane, fx.make_layout((warp_size // mfma_n, mfma_n), (mfma_n, 1)) + ) + lane_div_n = fx.Int32(fx.get(lane_split, 0)) + lane_mod_n = fx.Int32(fx.get(lane_split, 1)) + return wave_id, lane, lane_div_n, lane_mod_n + + +def grouped_loader(t, dim: int, g: int): + """Return a loader that reads a contiguous g-wide vector from a jagged 3D + tensor t[row, head, :], grouping the trailing dim into (dim/g, g) via a + layout so the group index selects the vector. + + The row*row_stride term is carried by the tensor arg's own i64-strided + layout (the leading index is i64), so the base address cannot overflow on + large packed tensors; the g-wide in-row vector load is a small, coalesced + access over the < int32 in-row span. + """ + in_row = fx.make_layout((dim // g, g), (g, 1)) + + def load(row_i64, head_val, colgrp): + sub = t[row_i64, head_val, None] + return fx.make_view(fx.get_iter(sub), in_row)[colgrp, None].load() + + return load + + +def swz_col(tile_row, col, swz_rows: int, swz_shift: int): + """XOR swizzle of an LDS column by the tile row (period swz_rows, shift + swz_shift). Shared by the forward K tile and the backward streamed Q/K LDS + tiles so the DMA global-fetch column and the LDS read column stay in sync. + """ + return col ^ ((tile_row & fx.Int32(swz_rows - 1)) << fx.Int32(swz_shift)) diff --git a/aiter/ops/flydsl/kernels/hstu_attention_fwd.py b/aiter/ops/flydsl/kernels/hstu_attention_fwd.py new file mode 100644 index 00000000000..c3bcc627328 --- /dev/null +++ b/aiter/ops/flydsl/kernels/hstu_attention_fwd.py @@ -0,0 +1,872 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""hstu_attention_fwd - FlyDSL kernel + +out_i = (1/N) * sum_j valid(i,j) * silu(alpha * q_i * k_j^T) * v_j + +HSTU forward for windowed symmetric heads and selected full-causal heads. GEMM2 consumes P as +operand A and row-major V as operand B, so V stages without a transposed LDS scatter. K stages by +dword DMA; V loads to registers, waits with counted vmcnt, and then publishes to LDS for GEMM2. +Row coordinates are cast to i64 so view address arithmetic cannot overflow on large packed tensors. + +Inputs/outputs: + - q,k (L, H, attn_dim); v,out (L, H, hidden_dim): packed jagged, `dtype`. Passed in their + native rank-3 layout (q/k/v must be contiguous); no host-side flatten to (L, H*dim). + - seq_offsets (Z+1) i32, num_targets (Z) i32. + +Paths: + - windowed (max_attn_len > 0): sliding-window lower bound skips fully-masked low KV tiles. + - full-causal (max_attn_len == 0): causal upper bound only. + - contextual (contextual_seq_len > 0): id shift+clamp, prefix-opener term, and the prefix + query block opens its KV range (upper→seq_len, lower→0) to see the whole prefix. + +Constraints: + - causal; contextual prefix requires causal. + - dtype in {f16, bf16}. + - attn_dim % 16 == 0, hidden_dim % 16 == 0; (batch*num_heads) % 8 == 0; tile divisibility is a + build-time contract. + - fast/unsafe FP math: not strict IEEE-754. +""" + +import functools +import math as host_math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir.dialects import fly +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl +from flydsl.expr.typing import Vector as Vec +from flydsl.runtime.device import get_rocm_arch +from flydsl.utils.smem_allocator import SMEM_CAPACITY_MAP + +from aiter.ops.flydsl.kernels.hstu_attention_common import ( + decode_lane, + grouped_loader, + swz_col, +) + + +def _dtype_to_elem_type(dtype_str: str): + if dtype_str == "f16": + return fx.Float16 + if dtype_str == "bf16": + return fx.BFloat16 + raise ValueError(f"unsupported dtype: {dtype_str!r} (expected 'f16' or 'bf16')") + + +# ---- Kernel Geometry Constants ---- + +WARP_SIZE = 64 +NUM_GRID_GROUPS = 8 # grid decoded group-major for locality +MFMA_M = 16 +MFMA_N = 16 +MFMA_K = 16 +MFMA_LANE_K = 4 +MFMA_ELEMS_PER_LANE = (MFMA_M * MFMA_N) // WARP_SIZE # 4 f32 per lane + +# K LDS swizzle and DMA granule are arch-conditional: bank conflict period +# differs between CDNA3 (128 B) and CDNA4 (256 B). Mask must stay < 64 elements so col ^ mask never +# leaves the row for any supported HEAD_DIM_K (including non-power-of-2 192). + + +def _arch_dma_params(arch: str | None = None): + """Arch-conditional (DMA_BYTES, DMA_ELEMS, K_SWZ_ROWS, K_SWZ_SHIFT). + + gfx942 (CDNA3): dword DMA (4 B/lane), 128-byte bank period -> (16, 2) covers a 64-element block. + gfx950 (CDNA4): dwordx4 DMA (16 B/lane), 256-byte bank period -> (8, 3) covers a 64-element block. + """ + if arch is None: + arch = get_rocm_arch() + if (arch or "").startswith("gfx942"): + dma_bytes = 4 # CDNA3 dword + k_swz_rows, k_swz_shift = 16, 2 + else: + dma_bytes = 16 # CDNA4 dwordx4 + k_swz_rows, k_swz_shift = 8, 3 + return dma_bytes, dma_bytes // 2, k_swz_rows, k_swz_shift + + +# V LDS is fully contiguous row-major [BLOCK_N, hidden_dim]: the dword buffer_load_lds scatters +# lanes contiguously from a wave-uniform base, so the LDS layout must match the global row-major +# fetch order. + + +@functools.lru_cache(maxsize=16384) +def lds_cap_bytes(arch: str | None = None) -> int: + default_lds_cap_bytes = 65536 + + if arch is None: + arch = get_rocm_arch() + return SMEM_CAPACITY_MAP.get(arch, default_lds_cap_bytes) + + +_LOG2E = host_math.log2(host_math.e) + + +def _waitcnt_vm_n(n: int): + """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7).""" + + # s_waitcnt vmcnt is split lo[3:0] @ bit 0, hi[5:4] @ bit 14. + # lgkmcnt(63) + expcnt(7) stay maximal so only vmcnt is constrained; V register loads remain + # outstanding while K DMA is drained. + vmcnt_lo_mask = 0xF + lgkmcnt_expcnt_base = 0x3F70 + vmcnt_hi_shift = 14 + vmcnt_hi_mask = 0x3 + val = ( + (n & vmcnt_lo_mask) + | lgkmcnt_expcnt_base + | (((n >> 4) & vmcnt_hi_mask) << vmcnt_hi_shift) + ) + rocdl.s_waitcnt(val) + + +def validate_hstu_attention_fwd( + num_heads: int, + head_dim: int, + hidden_dim: int, + batch: int, + causal: bool, + max_attn_len: int, + contextual_seq_len: int, + has_targets: bool, + alpha: float, + dtype_str: str, + max_seq_len: int, + *, + block_m: int, + block_n: int, + num_waves: int, + waves_per_eu: int, + arch: str | None = None, +) -> None: + if arch is None: + arch = get_rocm_arch() + if not arch.startswith("gfx942") and not arch.startswith("gfx950"): + raise ValueError( + f"hstu_attention_fwd unsupported arch: {arch!r} (expected 'gfx942' or 'gfx950')" + ) + + if dtype_str not in {"f16", "bf16"}: + raise ValueError(f"unsupported dtype: {dtype_str!r} (expected 'f16' or 'bf16')") + # `causal` is always True here; it is kept as an explicit parameter only for + # call-signature symmetry with the HSTU backward kernels (which share the same + # public wrapper). Non-causal HSTU is not a valid configuration. + if not causal: + raise ValueError("hstu_attention_fwd only supports causal attention") + if contextual_seq_len < 0: + raise ValueError( + f"contextual_seq_len must be non-negative, got {contextual_seq_len}" + ) + if max_attn_len < 0: + raise ValueError(f"max_attn_len must be non-negative, got {max_attn_len}") + if batch <= 0: + raise ValueError(f"batch must be positive, got {batch}") + if num_heads <= 0: + raise ValueError(f"num_heads must be positive, got {num_heads}") + if max_seq_len <= 0: + raise ValueError(f"max_seq_len must be positive, got {max_seq_len}") + if not host_math.isfinite(alpha): + raise ValueError(f"alpha must be finite, got {alpha}") + if head_dim <= 0 or head_dim % MFMA_K != 0: + raise ValueError( + f"head_dim must be positive and a multiple of MFMA_K={MFMA_K}, got {head_dim}" + ) + if hidden_dim <= 0 or hidden_dim % MFMA_M != 0: + raise ValueError( + f"hidden_dim must be positive and a multiple of MFMA_M={MFMA_M}, got {hidden_dim}" + ) + if (batch * num_heads) % NUM_GRID_GROUPS != 0: + raise ValueError( + f"require (batch*num_heads) % {NUM_GRID_GROUPS} == 0, got {batch * num_heads}" + ) + + if block_m <= 0: + raise ValueError(f"block_m must be positive, got {block_m}") + if block_n <= 0: + raise ValueError(f"block_n must be positive, got {block_n}") + if num_waves <= 0: + raise ValueError(f"num_waves must be positive, got {num_waves}") + if waves_per_eu < 0: + raise ValueError(f"waves_per_eu must be non-negative, got {waves_per_eu}") + if block_m % (num_waves * MFMA_M) != 0: + raise ValueError( + f"block_m {block_m} must be a multiple of num_waves*MFMA_M ({num_waves * MFMA_M})" + ) + if block_n % MFMA_M != 0: + raise ValueError(f"block_n {block_n} must be a multiple of MFMA_M={MFMA_M}") + + _, dma_elems, _, _ = _arch_dma_params() + block_threads = num_waves * WARP_SIZE + elems_per_dma_pass = block_threads * dma_elems + head_dim_k = ((head_dim + 63) // 64) * 64 + if (block_n * head_dim_k) % elems_per_dma_pass != 0: + raise ValueError("K DMA tile does not divide the dword DMA pass evenly") + if (block_n * hidden_dim) % elems_per_dma_pass != 0: + raise ValueError("V DMA tile does not divide the dword DMA pass evenly") + + vec_v = ( + 8 + if (hidden_dim % 8 == 0 and (block_n * hidden_dim) % (block_threads * 8) == 0) + else dma_elems + ) + threads_per_row_v = hidden_dim // vec_v + if block_threads % threads_per_row_v != 0: + raise ValueError( + f"block_threads={block_threads} must be divisible by threads_per_row_v={threads_per_row_v}" + ) + rows_per_batch_v = block_threads // threads_per_row_v + if not (block_n % rows_per_batch_v == 0 or rows_per_batch_v > block_n): + raise ValueError( + f"rows_per_batch_v={rows_per_batch_v} must divide block_n={block_n}, unless rows_per_batch_v > block_n" + ) + + lds_cap = lds_cap_bytes() + lds_bytes = block_n * head_dim_k * 2 + block_n * hidden_dim * 2 + if lds_bytes > lds_cap: + raise ValueError(f"LDS tile {lds_bytes} B exceeds the {lds_cap} B budget") + + +@functools.lru_cache(maxsize=16384) +def build_hstu_attention_fwd( + num_heads: int, + head_dim: int, + hidden_dim: int, + batch: int, + causal: bool, + max_attn_len: int, + contextual_seq_len: int, + has_targets: bool, + alpha: float, + dtype_str: str, + max_seq_len: int, + *, + # keyword-only tunables. These defaults are correctness-safe, not perf-tuned: + # production selects tuned (block_m, block_n, num_waves, waves_per_eu) per shape + # from the per-arch tuned CSV via the kernel wrapper. + block_m: int = 64, + block_n: int = 16, + num_waves: int = 2, + waves_per_eu: int = 0, +): + validate_hstu_attention_fwd( + num_heads, + head_dim, + hidden_dim, + batch, + causal, + max_attn_len, + contextual_seq_len, + has_targets, + alpha, + dtype_str, + max_seq_len, + block_m=block_m, + block_n=block_n, + num_waves=num_waves, + waves_per_eu=waves_per_eu, + ) + + BLOCK_M = block_m + BLOCK_N = block_n + NUM_WAVES = num_waves + BLOCK_THREADS = NUM_WAVES * WARP_SIZE + ROWS_PER_WAVE = BLOCK_M // NUM_WAVES + Q_SUBTILES = ROWS_PER_WAVE // MFMA_M # query sub-tiles per wave + KV_SUBTILES = BLOCK_N // MFMA_N # KV sub-tiles + WAVES_PER_EU = waves_per_eu + + assert num_waves > 0 and block_m % (num_waves * MFMA_M) == 0 + assert block_n % MFMA_M == 0 + assert head_dim % MFMA_K == 0 + assert hidden_dim % MFMA_M == 0 + assert (batch * num_heads) % NUM_GRID_GROUPS == 0 + + # Arch-conditional DMA + K LDS swizzle period + DMA_BYTES, DMA_ELEMS, K_SWZ_ROWS, K_SWZ_SHIFT = _arch_dma_params() + + elem_dtype = _dtype_to_elem_type(dtype_str) + is_bf16 = dtype_str == "bf16" + has_window = max_attn_len > 0 + has_contextual = contextual_seq_len > 0 + + K_STEPS = head_dim // MFMA_K # real 16-wide contraction steps (Q side) + # The XOR swizzle formula (k_swz_col) has period 64 elements. When K_STRIDE < 64 + # (head_dim % 64 != 0) the swizzled column lands outside [0, K_STRIDE), producing wrong LDS + # addresses on both the DMA write and the LDS read. GEMM1 is unaffected (Q is register-resident), + # but GEMM2's A-operand (P packs read back through the swizzled K LDS) gets corrupted data. + # Round head_dim up to a 64-aligned stride so the swizzle stays in-row by construction. + # Extra K columns over-fetch out-of-range data (buffer bounds -> 0) and pair with a zero Q + # operand, contributing nothing. The aligned case leaves HEAD_DIM_K == head_dim. + HEAD_DIM_K = ((head_dim + 63) // 64) * 64 + K_STEPS_K = HEAD_DIM_K // MFMA_K # padded steps (K side); always a multiple of 4 + D_CHUNKS = hidden_dim // MFMA_M # O accumulator / GEMM2 chunks + + num_q_tiles = (max_seq_len + BLOCK_M - 1) // BLOCK_M + hz_per_group = (batch * num_heads) // NUM_GRID_GROUPS + + stride_qk_n = num_heads * head_dim + + K_STRIDE = HEAD_DIM_K # no PAD_K - XOR swizzle replaces it; 64-aligned so the swizzle stays in-row + + # V LDS: row-major V[kv, d], no row pad. The contiguous lane scatter matches the + # natural global layout. + V_STRIDE = hidden_dim + + # K DMA tiling (dword passes). + k_tile_elems = BLOCK_N * K_STRIDE + elems_per_dma_pass = BLOCK_THREADS * DMA_ELEMS + assert k_tile_elems % elems_per_dma_pass == 0 + NUM_DMA_K = k_tile_elems // elems_per_dma_pass + PAIRS_PER_ROW_K = K_STRIDE // DMA_ELEMS + + # V tile divisibility (the DMA pass also gates the register-load tiling below). + v_tile_elems = BLOCK_N * hidden_dim + assert v_tile_elems % elems_per_dma_pass == 0 + + # V register-prefetch tiling. Each lane reads VEC_V contiguous d elements + # of one V row, coalesced (dwordx{VEC_V//2}). One pass = ROWS_PER_BATCH_V rows; NUM_BATCHES_V + # passes cover the BLOCK_N×hidden_dim tile. Coalesced AND overlappable (register dest, not + # buffer_load_lds), so the load can stay in flight across GEMM1 under a counted vmcnt. + VEC_V = ( + 8 + if (hidden_dim % 8 == 0 and (BLOCK_N * hidden_dim) % (BLOCK_THREADS * 8) == 0) + else DMA_ELEMS + ) + THREADS_PER_ROW_V = hidden_dim // VEC_V + assert BLOCK_THREADS % THREADS_PER_ROW_V == 0 + ROWS_PER_BATCH_V = BLOCK_THREADS // THREADS_PER_ROW_V + assert BLOCK_N % ROWS_PER_BATCH_V == 0 or ROWS_PER_BATCH_V > BLOCK_N + NUM_BATCHES_V = max(1, BLOCK_N // ROWS_PER_BATCH_V) + V_NEEDS_GUARD = ROWS_PER_BATCH_V > BLOCK_N + + # LDS map: [K row-major tile][V row-major tile]. K is XOR-swizzled by column; V stays + # natural [kv, d] so GEMM2 consumes it as operand B without a transpose scatter. Each field is + # a 16B-aligned fx.Array; SharedAllocator sizes the LDS global (no manual finalize/get_base). + @fx.struct + class SharedStorage: + k: fx.Array[elem_dtype, BLOCK_N * K_STRIDE, 16] + v: fx.Array[elem_dtype, BLOCK_N * V_STRIDE, 16] + + # ---- Device Kernel ---- + @flyc.kernel(known_block_size=[BLOCK_THREADS, 1, 1]) + def hstu_attention_fwd( + q: fx.Tensor, + k: fx.Tensor, + v: fx.Tensor, + seq_offsets: fx.Tensor, + num_targets: fx.Tensor, + out: fx.Tensor, + ) -> None: + elem_type = elem_dtype.ir_type + compute_type = fx.Float32.ir_type + v4f32_type = Vec.make_type(MFMA_ELEMS_PER_LANE, fx.Float32) + c_zero_mfma_pack = Vec.filled(MFMA_LANE_K, 0.0, elem_dtype).ir_value() + + # ---- MMA atom (layout algebra): one 16x16x16 f16/bf16 accumulate ---- + _mma_atom = fx.make_mma_atom(fx.rocdl.MFMA(MFMA_M, MFMA_M, MFMA_K, elem_dtype)) + + def mfma_acc(a_pack, b_pack, c): + """MFMA accumulate through the layout MMA atom.""" + return fly.mma_atom_call_ssa([v4f32_type], _mma_atom, a_pack, b_pack, c) + + # ---- Thread / lane indices (layout-algebra decode) ---- + tid = fx.Int32(gpu.thread_idx.x) + wave_id, _lane, lane_div_16, lane_mod_16 = decode_lane( + tid, NUM_WAVES, WARP_SIZE, MFMA_N + ) + + # ---- Group-major grid decode -> (batch_idx, head_idx, q_tile_idx) ---- + block_id = fx.Int32(gpu.block_idx.x) + grid_group = block_id % fx.Int32(NUM_GRID_GROUPS) + pos_in_group = block_id // fx.Int32(NUM_GRID_GROUPS) + local_hz_idx = pos_in_group // fx.Int32(num_q_tiles) + q_tile_idx = pos_in_group % fx.Int32(num_q_tiles) + hz_idx = grid_group * fx.Int32(hz_per_group) + local_hz_idx + batch_idx = hz_idx // fx.Int32(num_heads) + head_idx = hz_idx % fx.Int32(num_heads) + + # ---- Sequence bounds + id clamps (target tail) ---- + seq_start = fx.Int32(seq_offsets[batch_idx]) + seq_len = fx.Int32(seq_offsets[batch_idx + fx.Int32(1)]) - seq_start + + num_target = fx.Int32(0) + if has_targets: + num_target = fx.Int32(num_targets[batch_idx]) + + # Contextual shifts max_id BEFORE the target-tail clamp (oracle _valid_attn_mask order). + max_id = seq_len + if has_contextual: + max_id = seq_len - fx.Int32(contextual_seq_len) + fx.Int32(1) + if has_targets: + max_id = (num_target > fx.Int32(0)).select(max_id - num_target, max_id) + + # ---- Global tensor views (layout-algebra access; addresses carried by make_layout) ---- + # grouped_loader (shared) groups rows into g-wide coordinate slices for coalesced vector + # loads. The row*row_stride term is carried in 64-bit by the tensor arg's own i64-strided + # layout, so the base address cannot overflow on large packed tensors. + q_load = grouped_loader(q, head_dim, MFMA_LANE_K) + v_load = grouped_loader( + v, hidden_dim, VEC_V + ) # V register-prefetch path (coalesced global -> regs) + + q_head_offset = head_idx * fx.Int32(head_dim) + + # ---- K DMA base (per-(seq,head) byte offset folded into the rebased descriptor) ---- + # V uses a register-prefetch path (async_load_v_regs via v_load), not a buffer resource. + k_base_byte_offset = ( + fx.Int64(seq_start) * fx.Int64(stride_qk_n) + fx.Int64(q_head_offset) + ) * fx.Int64(2) + + # LDS as shape-carried 2D views: [BLOCK_N, stride] so an LDS access is Vec.load/store on + # (row, col) with the stride carried by the memref layout (no manual row*stride+col). The + # buffer is still contiguous row-major (K_STRIDE has no pad, V unpadded), so the DMA base + # extraction and contiguous buffer_load_lds fetch order are preserved. + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + # Shape-carried LDS views. Grouped by the access vector width (K packs MFMA_LANE_K, V publish + # VEC_V) so a load/store is view[row, colgrp, None].load()/.store(); the trailing group axis + # carries the stride (no manual row*stride+col). V reads gather single [kv, d] scalars. + k_read_view = lds.k.view( + fx.make_layout( + (BLOCK_N, K_STRIDE // MFMA_LANE_K, MFMA_LANE_K), + (K_STRIDE, MFMA_LANE_K, 1), + ) + ) + # One V view (grouped by VEC_V) serves both the coalesced publish store and the scalar GEMM2 + # gather, so only a single LDS base pointer stays live (register-pressure sensitive). + v_view = lds.v.view( + fx.make_layout((BLOCK_N, V_STRIDE // VEC_V, VEC_V), (V_STRIDE, VEC_V, 1)) + ) + k_lds_byte_base = buffer_ops.extract_base_index(k_read_view, address_space=3) + + # ── Copy-atom global->LDS DMA (buffer_load_lds via fx.copy) ── + # A BufferCopyLDS atom drives the buffer_load_lds instruction through the FlyDSL copy-atom. + # The atom hardcodes the cache-policy/aux operand to 0. + _dma_atom = fx.make_copy_atom( + fx.rocdl.BufferCopyLDS(DMA_BYTES * 8), DMA_BYTES * 8 + ) + _lds_ptr_ty = fx.PointerType.get(elem_type, 2, DMA_BYTES) + + def _rebased_buffer_div(base_iter, byte_off, n_elems): + # Fold the (large) seq/head base into the 48-bit descriptor base so the per-lane element + # index stays a small 32-bit voffset; max_size (0xFFFFFFFF) num-records. Public + # make_buffer_tensor assembles the V# (buffer flags + descriptor); we only pre-shift the + # base pointer, mirroring the buffer_load_lds path in kernels/gemm/preshuffle_gemm.py. + base_i64 = fx.Int64(fx.ptrtoint(base_iter)) + shifted = fx.inttoptr(base_iter.type, base_i64 + fx.Int64(byte_off)) + g_flat = fx.rocdl.make_buffer_tensor( + fx.Tensor( + fx.make_view( + shifted, fx.make_layout(fx.Int32(n_elems), fx.Int32(1)) + ) + ) + ) + return fx.logical_divide(g_flat, fx.make_layout(1, 1)) + + k_div = _rebased_buffer_div( + fx.get_iter(k), k_base_byte_offset, max_seq_len * stride_qk_n + ) + + # Single source of truth for the K LDS swizzle (shared): XOR the column with the tile row's + # low bits. Shared by the DMA global-fetch column and the LDS read column. + def k_swz_col(tile_row, col): + return swz_col(tile_row, col, K_SWZ_ROWS, K_SWZ_SHIFT) + + q_wave_base = q_tile_idx * fx.Int32(BLOCK_M) + wave_id * fx.Int32(ROWS_PER_WAVE) + + # ---- Q rows / bounds per query sub-tile ---- + q_rows = [] + q_in_bounds = [] + for qg in range_constexpr(Q_SUBTILES): + local = q_wave_base + fx.Int32(qg * MFMA_M) + lane_mod_16 + q_rows.append(local) + q_in_bounds.append(local < seq_len) + + # ---- Q B-operand packs (register-resident, per query sub-tile) ---- + q_packs = [] # q_packs[ks][qg] + for ks in range_constexpr(K_STEPS): + q_col = fx.Int32(ks * MFMA_K) + lane_div_16 * fx.Int32( + MFMA_LANE_K + ) # column within the head + per_qg = [] + for qg in range_constexpr(Q_SUBTILES): + safe = q_in_bounds[qg].select(seq_start + q_rows[qg], seq_start) + raw = q_load( + fx.Int64(safe), head_idx, q_col // fx.Int32(MFMA_LANE_K) + ).ir_value() + per_qg.append(q_in_bounds[qg].select(raw, c_zero_mfma_pack)) + q_packs.append(per_qg) + + # ---- Score-gate helpers ---- + c_alpha = fx.Float32(alpha) + c_inv_n = fx.Float32(1.0 / max_seq_len) + c_neg_log2e = fx.Float32(-_LOG2E) + c_one_f = fx.Float32(1.0) + c_zero_f = fx.Float32(0.0) + + def silu_scale_batch(s_list): + """Saturating silu(alpha*s) for a list of scores, stage-batched for ILP. + 1/N is hoisted to the O epilogue. + + The fastmath context gives every add/mul the `fast` flag; exp2 and rcp + both go through the rocdl builders (rocdl.exp2 -> v_exp_f32 approximate + hardware op; math.exp2 would lower to a slower expansion).""" + with arith.fastmath(arith.FastMathFlags.fast): + sc = [s * c_alpha for s in s_list] + tt = [s * c_neg_log2e for s in sc] + emu = [fx.Float32(rocdl.exp2(compute_type, t.ir_value())) for t in tt] + den = [c_one_f + e for e in emu] + sig = [fx.Float32(rocdl.rcp(compute_type, d)) for d in den] + return [sc[i] * sig[i] for i in range(len(s_list))] + + def to_id(x): + """Raw position -> masked id (contextual prefix shift, then target-tail clamp).""" + xid = x + if has_contextual: + xid = xid - fx.Int32(contextual_seq_len - 1) + xid = (xid < fx.Int32(0)).select(fx.Int32(0), xid) + if has_targets: + xid = (xid > max_id).select(max_id, xid) + return xid + + def pack_p(vals): + """Pack 4 f32 scores into a bf16/f16 MFMA pack (the GEMM2 A-operand fragment).""" + if is_bf16: + c16 = fx.Int32(16) + cmask = fx.Int32(0xFFFF0000) + + def bf16_pair(lo_f32, hi_f32): + lo_i32 = fx.Float32(lo_f32).bitcast(fx.Int32) + hi_i32 = fx.Float32(hi_f32).bitcast(fx.Int32) + return (hi_i32 & cmask) | lo_i32.shrui(c16) + + pairs = [bf16_pair(vals[0], vals[1]), bf16_pair(vals[2], vals[3])] + return Vec.from_elements(pairs, fx.Int32).bitcast(elem_dtype).ir_value() + elems = [fx.Float32(v).to(elem_dtype) for v in vals] + return Vec.from_elements(elems, elem_dtype).ir_value() + + q_rows_i32 = q_rows + q_row_ids = [to_id(q_rows_i32[qg]) for qg in range_constexpr(Q_SUBTILES)] + + # ---- KV range: causal upper bound + active predicate ---- + q_start = q_tile_idx * fx.Int32(BLOCK_M) + q_end = q_start + fx.Int32(BLOCK_M) + base_upper = seq_len + if causal: + clamped = (q_end < seq_len).select(q_end, seq_len) + if has_contextual: + # The prefix block holds logical row id 0, which attends the whole contextual + # prefix (col_id < max_id) above its diagonal, so its KV range opens to seq_len. + # Other blocks are pure causal and their high tiles are fully masked. + ctx_block = q_start < fx.Int32(contextual_seq_len) + base_upper = ctx_block.select(seq_len, clamped) + else: + base_upper = clamped + active = q_start < seq_len + kv_upper = active.select(base_upper, fx.Int32(0)) + n_tiles = (kv_upper + fx.Int32(BLOCK_N - 1)) // fx.Int32(BLOCK_N) + + # ---- Sliding-window lower bound: skip fully-masked low KV tiles ---- + kv_tile_start = fx.Int32(0) + if has_window: + eff_q_low = (q_start < max_id).select(q_start, max_id) + kv_lower = eff_q_low - fx.Int32(max_attn_len) + kv_lower = (kv_lower > fx.Int32(0)).select(kv_lower, fx.Int32(0)) + win_tile_start = kv_lower // fx.Int32(BLOCK_N) + if has_contextual: + # The prefix block must walk KV from 0 to see the prefix; the window lower bound + # would otherwise skip the low tiles the prefix opener needs. + ctx_prefix_block = q_start < fx.Int32(contextual_seq_len) + kv_tile_start = ctx_prefix_block.select(fx.Int32(0), win_tile_start) + else: + kv_tile_start = win_tile_start + + N_ACC = D_CHUNKS * Q_SUBTILES + c_zero_v4f32 = Vec.filled(MFMA_ELEMS_PER_LANE, 0.0, fx.Float32).ir_value() + + # ---- K DMA: global -> LDS (dword, swizzled global fetch) ---- + c_dma_elems = fx.Int32(DMA_ELEMS) + c_pairs_per_row_k = fx.Int32(PAIRS_PER_ROW_K) + + wave_lds_base_k = fx.Int64(k_lds_byte_base) + fx.Int64(wave_id) * fx.Int64( + WARP_SIZE * DMA_BYTES + ) + wave_lds_lane0_k = rocdl.readfirstlane(fx.Int64.ir_type, wave_lds_base_k) + k_dma_rows = [] + k_dma_gcols = [] + for d in range_constexpr(NUM_DMA_K): + pair = tid + fx.Int32(d * BLOCK_THREADS) + row = pair // c_pairs_per_row_k + col_pair = pair % c_pairs_per_row_k + col = col_pair * c_dma_elems + k_dma_rows.append(row) + k_dma_gcols.append(k_swz_col(row, col)) + + c_stride_qk_n = fx.Int32(stride_qk_n) + + def async_load_k(kv_start): + """Issue the async K[kv_start] dword DMA passes, global->LDS (swizzled), via the + BufferCopyLDS copy-atom.""" + for d in range_constexpr(NUM_DMA_K): + row = k_dma_rows[d] + in_bounds = (kv_start + row) < seq_len + local_tok = in_bounds.select(kv_start + row, fx.Int32(0)) + src_elem = local_tok * c_stride_qk_n + k_dma_gcols[d] + lds_byte = fx.Int32( + wave_lds_lane0_k + fx.Int64(d * BLOCK_THREADS * DMA_BYTES) + ) + dst = fx.make_view( + fx.inttoptr(_lds_ptr_ty, lds_byte), fx.make_layout(1, 1) + ) + src = fx.slice(k_div, (None, fx.Int32(src_elem))) + fx.copy(_dma_atom, src, dst) + + # ---- V register prefetch: coalesced global -> registers (non-blocking) ---- + # Lane (tid) owns row = tid//THREADS_PER_ROW_V (+ batch*ROWS_PER_BATCH_V), col base = + # (tid%THREADS_PER_ROW_V)*VEC_V. Reads VEC_V contiguous d elements. The load is issued but + # NOT waited: GEMM1 runs while it is in flight, the wait deferred to a counted vmcnt(0). + v_load_row_in_batch = tid // fx.Int32(THREADS_PER_ROW_V) + v_load_lane_in_row = tid % fx.Int32(THREADS_PER_ROW_V) + v_load_col = v_load_lane_in_row * fx.Int32(VEC_V) # column within the head + + def async_load_v_regs(kv_start): + """Issue coalesced V[kv_start] global loads to registers; return the vecs (non-blocking).""" + vecs = [] + for b in range_constexpr(NUM_BATCHES_V): + row = v_load_row_in_batch + fx.Int32(b * ROWS_PER_BATCH_V) + tok = kv_start + row + in_bounds = tok < seq_len + if V_NEEDS_GUARD: + in_bounds = in_bounds & (row < fx.Int32(BLOCK_N)) + safe_tok = in_bounds.select(seq_start + tok, seq_start) + raw = v_load( + fx.Int64(safe_tok), head_idx, v_load_col // fx.Int32(VEC_V) + ).ir_value() + vecs.append( + in_bounds.select(raw, Vec.filled(VEC_V, 0.0, elem_dtype).ir_value()) + ) + return vecs + + def store_v_regs_to_lds(vecs): + """Write prefetched V vecs to LDS row-major V[kv, d] (the GEMM2 B layout), 2D-indexed.""" + for b in range_constexpr(NUM_BATCHES_V): + row = v_load_row_in_batch + fx.Int32(b * ROWS_PER_BATCH_V) + v_view[row, v_load_col // fx.Int32(VEC_V), None].store(Vec(vecs[b])) + + # ==== GEMM1: Q·K^T -> P (P fragment already in GEMM2 A-operand layout) ==== + def read_k_packs(ng): + """LDS-read K A-operand packs for sub-tile ng (2D-indexed; col via the shared swizzle).""" + k_row = fx.Int32(ng * MFMA_M) + lane_mod_16 + packs = [] + for ks in range_constexpr(K_STEPS_K): + k_col = fx.Int32(ks * MFMA_K) + lane_div_16 * fx.Int32(MFMA_LANE_K) + # swz_col is MFMA_LANE_K-aligned (k_col and the swizzle shift are both >= that + # granularity), so //MFMA_LANE_K selects the packed-vector group. + packs.append( + k_read_view[ + k_row, k_swz_col(k_row, k_col) // fx.Int32(MFMA_LANE_K), None + ].load() + ) + return packs + + def compute_p_tile(kv_start, k_packs_by_ng): + """Q·K^T MFMA, apply the mask, silu-gate -> P packs for one KV tile. + + k_packs_by_ng: LDS-read A-operand packs [ng][ks] from `read_k_packs`.""" + p_packs = [ + [None for _ in range_constexpr(Q_SUBTILES)] + for _ in range_constexpr(KV_SUBTILES) + ] + for ng in range_constexpr(KV_SUBTILES): + k_packs = [ + Vec(k_packs_by_ng[ng][ks]) for ks in range_constexpr(K_STEPS_K) + ] + kv_base = ( + kv_start + + fx.Int32(ng * MFMA_M) + + lane_div_16 * fx.Int32(MFMA_LANE_K) + ) + col_raw = [ + kv_base + fx.Int32(i) for i in range_constexpr(MFMA_ELEMS_PER_LANE) + ] + col_in_seq = [ + col_raw[i] < seq_len for i in range_constexpr(MFMA_ELEMS_PER_LANE) + ] + col_id = [ + to_id(col_raw[i]) for i in range_constexpr(MFMA_ELEMS_PER_LANE) + ] + for qg in range_constexpr(Q_SUBTILES): + cur = Vec.filled(MFMA_ELEMS_PER_LANE, 0.0, fx.Float32).ir_value() + for ks in range_constexpr(K_STEPS_K): + q_op = q_packs[ks][qg] if ks < K_STEPS else c_zero_mfma_pack + cur = mfma_acc(k_packs[ks].ir_value(), q_op, cur) + s_vals = [Vec(cur)[i] for i in range_constexpr(MFMA_ELEMS_PER_LANE)] + + def keep_col( + i, + qg=qg, + col_id=col_id, + col_raw=col_raw, + col_in_seq=col_in_seq, + ): + """causal · window · contextual · target mask for (qg, col i).""" + dist = q_row_ids[qg] - col_id[i] + keep = (q_rows_i32[qg] == col_raw[i]) | (dist > fx.Int32(0)) + if has_window: + keep = keep & (dist <= fx.Int32(max_attn_len)) + if has_contextual: + # Prefix opener: logical row 0 attends the whole contextual prefix. + ctx = (q_row_ids[qg] == fx.Int32(0)) & (col_id[i] < max_id) + keep = keep | ctx + keep = keep & col_in_seq[i] + return keep + + s_vals = [ + keep_col(i).select(s_vals[i], c_zero_f) + for i in range_constexpr(MFMA_ELEMS_PER_LANE) + ] + p_packs[ng][qg] = pack_p(silu_scale_batch(s_vals)) + return p_packs + + # ==== GEMM2: P·V -> O (V as natural-layout operand B; P as operand A) ==== + def accum_o_tile(o_acc, p_packs): + """O[m,d] += P[m,n]·V[n,d]. A = P (GEMM1 frag, M=query K=kv), B = V[kv,d] natural.""" + for c in range_constexpr(D_CHUNKS): + # V operand B fragment for this d-chunk: tid%16 = d, and the 4 packed elements are + # stride-16 kv positions in row-major V[kv, d]: + # kv = ng*16 + 4*lane_div_16 + i; d = c*16 + lane_mod_16. + v_packs = [] + for ng in range_constexpr(KV_SUBTILES): + d_col = fx.Int32(c * MFMA_M) + lane_mod_16 # d index (N of GEMM2) + kv_lane = fx.Int32(ng * MFMA_M) + lane_div_16 * fx.Int32( + MFMA_LANE_K + ) # base kv + # b[i] = V[kv_lane + i, d_col]; scalar-gathered from the VEC_V-grouped V view + # (col d_col -> group d_col//VEC_V, lane d_col%VEC_V). + d_grp = d_col // fx.Int32(VEC_V) + d_lane = d_col % fx.Int32(VEC_V) + elems = [ + v_view[kv_lane + fx.Int32(i), d_grp, d_lane] + for i in range_constexpr(MFMA_LANE_K) + ] + v_packs.append(Vec.from_elements(elems, elem_dtype).ir_value()) + for qg in range_constexpr(Q_SUBTILES): + acc_off = c * Q_SUBTILES + qg + cur = o_acc[acc_off] + for ng in range_constexpr(KV_SUBTILES): + cur = mfma_acc(p_packs[ng][qg], v_packs[ng], cur) + o_acc[acc_off] = cur + return o_acc + + # ==== Main pipeline (single K/V LDS slot, V register-prefetch overlap) ==== + # V is prefetched by ordinary global loads, so the K publish barrier does not drain it. + # GEMM1 runs while V is in flight; vmcnt(0) drains V before the LDS publish for GEMM2. + # Only the O accumulators are loop-carried. + v_reg_outstanding = ( + NUM_BATCHES_V # V register loads kept in flight while waiting on K + ) + + def run_kv_tile(o_acc, kv_start): + """K staged global->LDS by swizzled DMA; V register-prefetched.""" + async_load_k(kv_start) # K[i] global -> LDS (GEMM1 A-operand) + v_vecs = async_load_v_regs(kv_start) # V[i] global -> regs, in flight + _waitcnt_vm_n( + v_reg_outstanding + ) # wait K[i] DMA only; V[i] stays outstanding + gpu.barrier() # K[i] LDS published (V[i] regs survive the barrier) + k_packs = [read_k_packs(ng) for ng in range_constexpr(KV_SUBTILES)] + p_packs = compute_p_tile( + kv_start, k_packs + ) # GEMM1(i) overlaps V[i] global load + _waitcnt_vm_n(0) # V[i] arrived + store_v_regs_to_lds(v_vecs) + rocdl.sched_dswr( + NUM_BATCHES_V + ) # keep the NUM_BATCHES_V V LDS writes grouped + gpu.barrier() # V[i] LDS published + o_acc = accum_o_tile(o_acc, p_packs) # GEMM2(i): O += P·V + # The next iteration's K publish barrier also fences GEMM2(i)'s V LDS reads before + # store_v(i+1). + return o_acc + + if active: + acc_init = [c_zero_v4f32 for _ in range(N_ACC)] + loop_results = acc_init + for kv_tile, it in range( + fx.Int64(kv_tile_start), fx.Int64(n_tiles), fx.Int64(1), init=acc_init + ): # ty: ignore + it_list = list(it) if isinstance(it, (list, tuple)) else [it] + o_acc = [it_list[i] for i in range(N_ACC)] + kv_start = fx.Int32(kv_tile) * fx.Int32(BLOCK_N) + o_acc = run_kv_tile(o_acc, kv_start) + loop_results = yield o_acc + + # ---- Epilogue: store O (1/N hoisted here) ---- + # GEMM2 writes a transposed C fragment: A=P, B=V, C[d, query]. + # tid%16 = d (N-dim), and (lane_div_16, e) = query (M-dim). The query row stored is + # therefore q_wave_base + qg*16 + lane_div_16*4 + e; the d column is c*16 + lane_mod_16. + results = ( + list(loop_results) + if isinstance(loop_results, (list, tuple)) + else [loop_results] + ) + with arith.fastmath(arith.FastMathFlags.fast): + for qg in range_constexpr(Q_SUBTILES): + q_row_base = ( + q_wave_base + + fx.Int32(qg * MFMA_M) + + lane_div_16 * fx.Int32(MFMA_LANE_K) + ) + for e in range_constexpr(MFMA_ELEMS_PER_LANE): + q_row_e = q_row_base + fx.Int32(e) + if q_row_e < seq_len: + for c in range_constexpr(D_CHUNKS): + ov = results[c * Q_SUBTILES + qg] + d_col = fx.Int32(c * MFMA_M) + lane_mod_16 + val = (Vec(ov)[e] * c_inv_n).to(elem_dtype) + out[fx.Int64(seq_start + q_row_e), head_idx, d_col] = ( + val + ) + + _hstu_compile_hints = { + "fast_fp_math": True, + "unsafe_fp_math": True, + } + + @flyc.jit + def launch_hstu_attention_fwd( + q: fx.Tensor, + k: fx.Tensor, + v: fx.Tensor, + seq_offsets: fx.Tensor, + num_targets: fx.Tensor, + out: fx.Tensor, + stream: fx.Stream, + ) -> None: + grid = num_q_tiles * batch * num_heads + hstu_attention_fwd( + q, + k, + v, + seq_offsets, + num_targets, + out, + value_attrs={ + "passthrough": [ + ["denormal-fp-math-f32", "preserve-sign,preserve-sign"], + ["no-nans-fp-math", "true"], + ["unsafe-fp-math", "true"], + ], + "rocdl.waves_per_eu": WAVES_PER_EU, + "rocdl.flat_work_group_size": f"{BLOCK_THREADS},{BLOCK_THREADS}", + }, + ).launch( + grid=grid, + block=BLOCK_THREADS, + smem=0, + stream=stream, + ) + + launch_hstu_attention_fwd.compile_hints = _hstu_compile_hints + return launch_hstu_attention_fwd diff --git a/op_tests/flydsl_tests/test_flydsl_hstu_attention.py b/op_tests/flydsl_tests/test_flydsl_hstu_attention.py new file mode 100644 index 00000000000..f632f2c093c --- /dev/null +++ b/op_tests/flydsl_tests/test_flydsl_hstu_attention.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import argparse +import csv +import itertools + +import pandas as pd +import pytest +import torch + +import aiter +import aiter.ops.flydsl.hstu_attention_kernels as hstu_kernels +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.hstu_attention_kernels import ( + _validate_inputs, + flydsl_hstu_attention_fwd, +) +from aiter.test_common import benchmark, checkAllclose, run_perftest + +# Every card this op is built/validated for; the kernel's validate() enforces the +# same set, but gate the whole perf sweep here so it is a clean no-op elsewhere. +SUPPORTED_GFX = ["gfx942", "gfx950"] + +# Module-level singleton so it is not constructed in a function default (ruff B008). +_DEFAULT_DEVICE = torch.device("cuda") + + +# --------------------------------------------------------------------------- # +# Input generation (shared with op_benchmarks/flydsl/bench_hstu_attn.py) +# --------------------------------------------------------------------------- # +def _generate_sparse_seq_len( + size: int, + max_seq_len: int, + sparsity: float, + device: torch.device, +) -> torch.Tensor: + torch.manual_seed(1) + + if sparsity == 0.0: + return torch.zeros(size=(size,), device=device, dtype=torch.int) + elif sparsity == 1.0: + return torch.ones(size=(size,), device=device, dtype=torch.int) * max_seq_len + elif sparsity >= 0.5: + min_seq_len = int((2 * sparsity - 1.0) * max_seq_len) + else: + min_seq_len = 0 + max_seq_len = int(2 * sparsity * max_seq_len) + + return torch.randint( + low=min_seq_len, + high=max_seq_len, + size=(size,), + device=device, + dtype=torch.int, + ) + + +def generate_hstu_attn_inputs( + batch_size: int, + max_seq_len: int, + sparsity: float, + heads: int, + attn_dim: int, + hidden_dim: int, + target_size: int, + dtype: torch.dtype, + device: torch.device = _DEFAULT_DEVICE, + seed: int = 1001, +): + """Build (q, k, v, seq_offsets, num_targets) for an HSTU attention problem. + + Self-contained (no triton-side dependencies) so both the tests here and the + FlyDSL benchmark can share a single input generator. + """ + torch.manual_seed(seed) # for reproducibility + + # sparsity-controlled jagged lengths (mean ~= sparsity * max_seq_len). The + # old power-law resampling step (apply_SL) is intentionally dropped: ported + # with alpha=0.2 it collapsed every sequence to ~2 tokens, which made the + # perf roofline meaningless. Realistic lengths here match the mvonstra + # recsys_harness convention used to validate this kernel. + lengths = _generate_sparse_seq_len(batch_size, max_seq_len, sparsity, device) + + num_targets = None + if target_size > 0: + num_targets = torch.randint( + 1, + target_size + 1, + (batch_size,), + device=lengths.device, + dtype=lengths.dtype, + ) + num_targets = torch.where(num_targets > lengths, lengths, num_targets) + + seq_offsets = torch.zeros((batch_size + 1,), dtype=torch.int64, device=device) + seq_offsets[1:] = torch.cumsum(lengths, dim=0) + total_len = int(seq_offsets[-1].item()) + + x = torch.empty( + (total_len, heads, attn_dim * 2 + hidden_dim), + dtype=dtype, + device=device, + ).uniform_(-0.01, 0.01) + q, k, v = torch.split(x, [attn_dim, attn_dim, hidden_dim], dim=-1) + + return q.contiguous(), k.contiguous(), v.contiguous(), seq_offsets, num_targets + + +# --------------------------------------------------------------------------- # +# Reference (oracle only — never timed, never in the summary table) +# --------------------------------------------------------------------------- # +def run_torch( + max_seq_len, + alpha, + q, + k, + v, + seq_offsets, + causal, + num_targets, + max_attn_len, + contextual_seq_len, +): + # The torch reference lives on the triton side; import it lazily so merely + # importing this module (e.g. from the benchmark) stays triton-free. + from op_tests.triton_tests.utils.hstu_attention_ref import torch_hstu_attention + + return torch_hstu_attention( + max_seq_len, + alpha, + q, + k, + v, + seq_offsets, + causal, + dropout_pr=0.0, + training=False, + num_targets=num_targets, + max_attn_len=max_attn_len, + contextual_seq_len=contextual_seq_len, + min_full_attn_seq_len=0, + ) + + +def _roofline(seq_offsets, heads, attn_dim, hidden_dim, elem_size): + """FLOPs and HBM bytes the op actually does, for TFLOPS / TB-s. + + Causal HSTU is two GEMMs per (query, key) pair: Q·K^T contracts attn_dim and + P·V contracts hidden_dim. Valid pairs ~ length^2/2 (lower-triangular) and each + pair is a multiply-add (2 flops), so the causal 1/2 and the MAC 2 cancel: + FLOPs ~ sum_b length^2 * (attn_dim + hidden_dim) * heads. + A fused kernel never materialises the L*L score matrix in HBM, so the traffic + is just q + k (attn_dim) read and v read + out written (hidden_dim): + bytes ~ sum_b length * heads * 2*(attn_dim + hidden_dim) * elem_size. + """ + lengths = (seq_offsets[1:] - seq_offsets[:-1]).to(torch.float64).cpu() + flops = float((lengths * lengths).sum()) * (attn_dim + hidden_dim) * heads + total_tokens = float(lengths.sum()) + nbytes = total_tokens * heads * 2 * (attn_dim + hidden_dim) * elem_size + return flops, nbytes + + +# --------------------------------------------------------------------------- # +# Perf + correctness sweep (aiter-op-test format) +# --------------------------------------------------------------------------- # +@benchmark() +def test_hstu_attention_fwd( + batch_size, + max_seq_len, + sparsity, + num_heads, + head_dim, + hidden_dim, + max_attn_len, + contextual_seq_len, + target_size, + dtype, +): + causal = True + alpha = 1.0 / head_dim * 10000 + + q, k, v, seq_offsets, num_targets = generate_hstu_attn_inputs( + batch_size=batch_size, + max_seq_len=max_seq_len, + sparsity=sparsity, + heads=num_heads, + attn_dim=head_dim, + hidden_dim=hidden_dim, + target_size=target_size, + dtype=dtype, + ) + + # torch reference is the oracle only — compute once, compare against it, never time it. + ref = run_torch( + max_seq_len, + alpha, + q, + k, + v, + seq_offsets, + causal, + num_targets, + max_attn_len, + contextual_seq_len, + ) + + candidates = { + "flydsl": lambda: flydsl_hstu_attention_fwd( + max_seq_len, + alpha, + q, + k, + v, + seq_offsets, + causal, + num_targets, + max_attn_len, + contextual_seq_len, + ), + } + + flops, nbytes = _roofline( + seq_offsets, num_heads, head_dim, hidden_dim, q.element_size() + ) + + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + # Output carries the 1/N normalisation; rescale both sides to O(1) so the + # bf16/f16 tolerance is meaningful, then compare in fp32. + err = checkAllclose( + (ref.to(dtypes.fp32) * max_seq_len), + (out.to(dtypes.fp32) * max_seq_len), + atol=1e-3, + rtol=0, + msg=f"{name}: hstu_attention_fwd", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +# @benchmark wraps the fn so its call args become table columns; it is driven by +# main() / the correctness test below, not collected as a bare pytest test. +test_hstu_attention_fwd.__test__ = False + + +# --------------------------------------------------------------------------- # +# Correctness (drives the @benchmark fn and asserts the kernel matches torch) +# --------------------------------------------------------------------------- # +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA device" +) + + +@requires_cuda +@pytest.mark.parametrize( + "batch_size,max_seq_len,sparsity," + "max_attn_len,contextual_seq_len,target_size," + "head_dim,hidden_dim", + [ + (256, 1024, 0.5, 0, 0, 0, 128, 128), + # target_size > 0 + (256, 1024, 0.5, 0, 0, 20, 128, 128), + # max_attn_len > 0 + (256, 1024, 0.5, 64, 0, 0, 128, 128), + # contextual_seq_len > 0 + (256, 1024, 0.5, 0, 64, 0, 128, 128), + # symmetric and dims %64 != 0 + (256, 1024, 0.5, 0, 0, 0, 96, 96), + # not symmetric + (256, 1024, 0.5, 0, 0, 0, 128, 64), + # not symmetric and dims %64 != 0 + (256, 1024, 0.5, 0, 0, 0, 96, 192), + ], +) +def test_flydsl_hstu_attention( + batch_size, + max_seq_len, + sparsity, + max_attn_len, + contextual_seq_len, + target_size, + head_dim, + hidden_dim, + num_heads=4, + dtype=torch.bfloat16, +): + if get_gfx() not in SUPPORTED_GFX: + pytest.skip(f"hstu_attention_fwd unsupported on {get_gfx()}") + + torch.cuda.empty_cache() + + ret = test_hstu_attention_fwd( + batch_size, + max_seq_len, + sparsity, + num_heads, + head_dim, + hidden_dim, + max_attn_len, + contextual_seq_len, + target_size, + dtype, + ) + assert ret["flydsl err"] == 0, ret + + +# --------------------------------------------------------------------------- # +# Input validation +# --------------------------------------------------------------------------- # +def _qkv(batch=2, tokens=8, heads=4, attn_dim=128, hidden_dim=128, device="cuda"): + q = torch.zeros((tokens, heads, attn_dim), dtype=torch.bfloat16, device=device) + k = torch.zeros_like(q) + v = torch.zeros((tokens, heads, hidden_dim), dtype=torch.bfloat16, device=device) + seq_offsets = torch.zeros(batch + 1, dtype=torch.int64, device=device) + return q, k, v, seq_offsets + + +@requires_cuda +def test_validate_inputs_ok(): + q, k, v, seq_offsets = _qkv(batch=2, heads=4, attn_dim=128, hidden_dim=96) + batch, num_heads, head_dim, hidden_dim, dtype_str = _validate_inputs( + q, k, v, seq_offsets, None + ) + + actual = (batch, num_heads, head_dim, hidden_dim, dtype_str) + expected = (2, 4, 128, 96, "bf16") + assert actual == expected + + +def test_validate_inputs_rejects_cpu_tensors(): + q, k, v, seq_offsets = _qkv(device="cpu") + with pytest.raises(ValueError): + _validate_inputs(q, k, v, seq_offsets, None) + + +@requires_cuda +def test_validate_inputs_rejects_qk_shape_mismatch(): + q, k, v, seq_offsets = _qkv() + k = torch.zeros( + (q.shape[0], q.shape[1], q.shape[2] + 8), dtype=q.dtype, device=q.device + ) + with pytest.raises(ValueError): + _validate_inputs(q, k, v, seq_offsets, None) + + +@requires_cuda +def test_validate_inputs_rejects_wrong_rank(): + q, k, v, seq_offsets = _qkv() + with pytest.raises(ValueError): + _validate_inputs(q.reshape(-1), k, v, seq_offsets, None) + + +@requires_cuda +def test_validate_inputs_rejects_dtype_mismatch(): + q, k, v, seq_offsets = _qkv() + v = v.to(torch.float16) + with pytest.raises(ValueError): + _validate_inputs(q, k, v, seq_offsets, None) + + +@requires_cuda +def test_validate_inputs_rejects_num_targets_length_mismatch(): + q, k, v, seq_offsets = _qkv(batch=2) + num_targets = torch.zeros(3, dtype=torch.int64, device=q.device) # != batch + with pytest.raises(ValueError): + _validate_inputs(q, k, v, seq_offsets, num_targets) + + +# --------------------------------------------------------------------------- # +# Tuned CSV loading +# --------------------------------------------------------------------------- # +def _row(**overrides) -> dict: + row = { + "arch": hstu_kernels._GPU_ARCH, + "dtype": "bf16", + "num_heads": 4, + "head_dim": 128, + "hidden_dim": 128, + "batch": 256, + "max_seq_len": 1024, + "has_window": "False", + "has_contextual": "False", + "has_targets": "False", + "duration": 1.0, + "block_m": 128, + "block_n": 64, + "num_waves": 4, + "waves_per_eu": 2, + } + row.update(overrides) + return row + + +def _write_csv(path, rows) -> str: + with open(path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=hstu_kernels._CSV_COLUMNS) + writer.writeheader() + for row in rows: + writer.writerow(row) + return str(path) + + +def test_tuned_csv_is_picked_up(tmp_path): + path = _write_csv(tmp_path / "tuned.csv", [_row()]) + + config_map = hstu_kernels._tuned_config_map(path) + + assert len(config_map) == 1 + (config,) = config_map.values() + assert config == { + "block_m": 128, + "block_n": 64, + "num_waves": 4, + "waves_per_eu": 2, + } + + +def test_tuned_csv_missing_file_returns_empty(tmp_path): + assert hstu_kernels._tuned_config_map(str(tmp_path / "does_not_exist.csv")) == {} + + +def test_tuned_csv_best_duration_wins(tmp_path): + path = _write_csv( + tmp_path / "tuned.csv", + [ + _row(duration=5.0, block_m=64), + _row(duration=1.0, block_m=256), + ], + ) + + config_map = hstu_kernels._tuned_config_map(path) + + (config,) = config_map.values() + assert config["block_m"] == 256 + + +# --------------------------------------------------------------------------- # +# Perf sweep entry point (markdown summary table) +# --------------------------------------------------------------------------- # +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning( + "hstu_attention_fwd unsupported on %s; skipping perf sweep", get_gfx() + ) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + choices=[dtypes.d_dtypes["bf16"], dtypes.d_dtypes["fp16"]], + nargs="*", + default="bf16,", + metavar="{bf16,fp16}", + help="Data type, e.g.: -d bf16", + ) + parser.add_argument( + "-b", + "--batch", + type=int, + nargs="*", + default=[128], + help="Batch size (num sequences). e.g.: -b 128", + ) + parser.add_argument( + "--seq", + type=int, + nargs="*", + default=[1024, 2048, 4096], + help="max_seq_len per sequence. e.g.: --seq 1024 2048", + ) + parser.add_argument( + "--sparsity", + type=float, + nargs="*", + default=[0.5], + help="Jagged sequence-length sparsity in [0,1]. e.g.: --sparsity 0.5", + ) + parser.add_argument( + "--heads", + type=int, + nargs="*", + default=[4], + help="num_heads. e.g.: --heads 4", + ) + parser.add_argument( + "-s", + "--dims", + type=dtypes.str2tuple, + nargs="*", + default=[(128, 128), (64, 64)], + help="(head_dim, hidden_dim) pairs. e.g.: -s 128,128 64,64", + ) + parser.add_argument( + "--max_attn_len", + type=int, + nargs="*", + default=[0], + help="Sliding-window length (0 = full causal). e.g.: --max_attn_len 0 256", + ) + parser.add_argument( + "--contextual_seq_len", + type=int, + nargs="*", + default=[0], + help="Contextual prefix length (0 = off). e.g.: --contextual_seq_len 0 64", + ) + parser.add_argument( + "--target_size", + type=int, + nargs="*", + default=[0], + help="Max target-tail size (0 = no targets). e.g.: --target_size 0 20", + ) + args = parser.parse_args() + + for dtype in args.dtype: + df = [] + for ( + batch, + seq, + sparsity, + heads, + (head_dim, hidden_dim), + max_attn_len, + contextual_seq_len, + target_size, + ) in itertools.product( + args.batch, + args.seq, + args.sparsity, + args.heads, + args.dims, + args.max_attn_len, + args.contextual_seq_len, + args.target_size, + ): + df.append( + test_hstu_attention_fwd( + batch, + seq, + sparsity, + heads, + head_dim, + hidden_dim, + max_attn_len, + contextual_seq_len, + target_size, + dtype, + ) + ) + df = pd.DataFrame(df) + aiter.logger.info( + "flydsl_hstu_attention_fwd summary (markdown):\n%s", + df.to_markdown(index=False), + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/op_benchmarks/flydsl/bench_hstu_attn.py b/op_tests/op_benchmarks/flydsl/bench_hstu_attn.py new file mode 100644 index 00000000000..8b83a68f4ee --- /dev/null +++ b/op_tests/op_benchmarks/flydsl/bench_hstu_attn.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import argparse +from dataclasses import dataclass + +import torch +import triton + +from aiter.ops.flydsl import flydsl_hstu_attention_fwd +from op_tests.flydsl_tests.test_flydsl_hstu_attention import ( + generate_hstu_attn_inputs, +) + + +# lower triangular mask, so no need to multiply by 2 for flops +def get_flops(seq_offsets: torch.Tensor, heads: int, attn_dim: int, hidden_dim: int): + total_flops = 0.0 + seq_num = seq_offsets.shape[0] - 1 + for i in range(seq_num): + length = seq_offsets[i + 1] - seq_offsets[i] + total_flops += length * length * (attn_dim + hidden_dim) * heads + return total_flops + + +def get_bytes( + seq_offsets: torch.Tensor, + heads: int, + attn_dim: int, + hidden_dim: int, + elem_size: int, +): + seq_num = seq_offsets.shape[0] - 1 + total_bytes = 0 + for i in range(seq_num): + length = seq_offsets[i + 1] - seq_offsets[i] + total_bytes += length * (attn_dim + length + hidden_dim) * heads * elem_size + return total_bytes + + +@dataclass +class FwdShape: + batch_size: int + max_seq_len: int + sparsity: float + dtype: str = "bf16" + num_heads: int = 4 + head_dim: int = 128 + hidden_dim: int = 128 + max_attn_len: int = 0 + contextual_seq_len: int = 0 + target_size: int = 20 + + def display_string(self) -> str: + parts: list[str] = [ + f"batch_size={self.batch_size}", + f"max_seq_len={self.max_seq_len}", + f"sparsity={self.sparsity}", + f"dtype={self.dtype}", + f"num_heads={self.num_heads}", + f"head_dim={self.head_dim}", + f"hidden_dim={self.hidden_dim}", + f"max_attn_len={self.max_attn_len}", + f"contextual_seq_len={self.contextual_seq_len}", + f"target_size={self.target_size}", + ] + return ",".join(parts) + + +_DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16} + +# Module-level singleton so it is not constructed in a function default (ruff B008). +_DEFAULT_DEVICE = torch.device("cuda") + + +def run_benchmark( + shape: FwdShape, + device: torch.device = _DEFAULT_DEVICE, + seed: int = 1001, +): + torch.cuda.empty_cache() + dtype = _DTYPES[shape.dtype] + + alpha = 1.0 / shape.head_dim * 10000 + causal = True + + q, k, v, seq_offsets, num_targets = generate_hstu_attn_inputs( + batch_size=shape.batch_size, + max_seq_len=shape.max_seq_len, + sparsity=shape.sparsity, + heads=shape.num_heads, + attn_dim=shape.head_dim, + hidden_dim=shape.hidden_dim, + target_size=shape.target_size, + dtype=dtype, + device=device, + seed=seed, + ) + + def flydsl_attn(): + return flydsl_hstu_attention_fwd( + shape.max_seq_len, + alpha, + q, + k, + v, + seq_offsets, + causal, + num_targets, + shape.max_attn_len, + shape.contextual_seq_len, + ) + + ms = triton.testing.do_bench(flydsl_attn, warmup=2000, rep=2000) + + flops = get_flops( + seq_offsets, + shape.num_heads, + shape.head_dim, + shape.hidden_dim, + ) + tflops = flops / ms / 1e9 + + elem_size = q.element_size() + bytes = get_bytes( + seq_offsets.cpu().numpy(), + shape.num_heads, + shape.head_dim, + shape.hidden_dim, + elem_size, + ) + bandwidth = bytes / (ms * 1e-3) * 1e-9 # GB/s + + return ms, tflops, bandwidth + + +def main(): + p = argparse.ArgumentParser(description="FlyDSL HSTU Attention benchmark") + p.add_argument("--batch_size", type=int, default=120) + p.add_argument("--max_seq_len", type=int, default=16384) + p.add_argument("--sparsity", type=float, default=0.475) + p.add_argument("--num_heads", type=int, default=4) + p.add_argument("--head_dim", type=int, default=64) + p.add_argument("--hidden_dim", type=int, default=0, help="defaults to --head_dim") + p.add_argument("--max_attn_len", type=int, default=0) + p.add_argument("--contextual_seq_len", type=int, default=0) + p.add_argument("--target_size", type=int, default=300) + p.add_argument("--dtype", type=str, default="bf16", choices=list(_DTYPES)) + args = p.parse_args() + + args.hidden_dim = args.hidden_dim or args.head_dim + + shape = FwdShape( + batch_size=args.batch_size, + max_seq_len=args.max_seq_len, + sparsity=args.sparsity, + dtype=args.dtype, + max_attn_len=args.max_attn_len, + contextual_seq_len=args.contextual_seq_len, + target_size=args.target_size, + head_dim=args.head_dim, + hidden_dim=args.hidden_dim, + num_heads=args.num_heads, + ) + ms, tflops, bandwidth = run_benchmark( + shape=shape, + device=torch.device("cuda"), + seed=1001, + ) + + print( + f"[FlyDSL HSTU Attention Forward] {shape.display_string()},ms={ms:.3f},tflops={tflops:.1f},gbps={bandwidth:.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/triton_tests/utils/hstu_attention_ref.py b/op_tests/triton_tests/utils/hstu_attention_ref.py index 2f69c1c893c..c0dd3e34129 100644 --- a/op_tests/triton_tests/utils/hstu_attention_ref.py +++ b/op_tests/triton_tests/utils/hstu_attention_ref.py @@ -1,5 +1,4 @@ import torch -from typing import Optional, Tuple import torch.nn.functional as F @@ -9,7 +8,7 @@ def _get_valid_attn_mask( causal: bool, N: int, seq_lengths: torch.Tensor, - num_targets: Optional[torch.Tensor] = None, + num_targets: torch.Tensor | None = None, max_attn_len: int = 0, contextual_seq_len: int = 0, min_full_attn_seq_len: int = 0, @@ -63,7 +62,7 @@ def jagged_to_padded_dense( q: torch.Tensor, offsets: torch.Tensor, max_seq_len: int, padding_value ): assert len(q.shape) == 2, "q needs to be 2-dim tensor" - L, D = q.shape + _L, D = q.shape B = offsets.shape[0] - 1 padded_shape = (B, max_seq_len, D) padded_q = torch.full(padded_shape, padding_value, dtype=q.dtype, device=q.device) @@ -91,21 +90,22 @@ def qkv_to_padded_dense( v: torch.Tensor, seq_offsets: torch.Tensor, N: int, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - L, H, D = q.shape +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + L, H, DQ = q.shape + DV = v.shape[2] padded_q = ( - pad_sequence(q.reshape(L, H * D), seq_offsets, N, 0.0) - .view(-1, N, H, D) + pad_sequence(q.reshape(L, H * DQ), seq_offsets, N, 0.0) + .view(-1, N, H, DQ) .transpose(1, 2) ) padded_k = ( - pad_sequence(k.reshape(L, H * D), seq_offsets, N, 0.0) - .view(-1, N, H, D) + pad_sequence(k.reshape(L, H * DQ), seq_offsets, N, 0.0) + .view(-1, N, H, DQ) .transpose(1, 2) ) padded_v = ( - pad_sequence(v.reshape(L, H * D), seq_offsets, N, 0.0) - .view(-1, N, H, D) + pad_sequence(v.reshape(L, H * DV), seq_offsets, N, 0.0) + .view(-1, N, H, DV) .transpose(1, 2) ) @@ -114,7 +114,7 @@ def qkv_to_padded_dense( # convert sequences from dense format to jagged format def dense_to_jagged(seq: torch.Tensor, offsets: torch.Tensor, L: int): - B, N, HV = seq.shape + B, _N, HV = seq.shape assert L == offsets[-1], f"jagged dim mismatch {offsets[-1]} != {L}!" out = torch.empty((L, HV), dtype=seq.dtype, device=seq.device) @@ -137,7 +137,7 @@ def torch_hstu_attention( causal: bool = True, dropout_pr: float = 0.0, training: bool = True, - num_targets: Optional[torch.Tensor] = None, + num_targets: torch.Tensor | None = None, max_attn_len: int = 0, contextual_seq_len: int = 0, min_full_attn_seq_len: int = 0,