diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index 2e95210388a..b56ed7df3a7 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -1493,6 +1493,22 @@ "f'{AITER_META_DIR}/hsa/codegen.py -m fmha_fwd_bf16_varlen --output_dir {{}}'" ] }, + "module_fmha_fwd_fp8_varlen_asm": { + "srcs": [ + "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu'" + ], + "flags_extra_cc": [ + "'-DENABLE_CK=0'" + ], + "flags_extra_hip": [], + "extra_ldflags": "None", + "extra_include": [], + "verbose": "False", + "torch_exclude": "True", + "blob_gen_cmd": [ + "f'{AITER_META_DIR}/hsa/codegen.py -m fmha_fwd_fp8_varlen --output_dir {{}}'" + ] + }, "module_fmha_fwd_mxfp8_asm": { "srcs": [ "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_fmha_fwd_mxfp8.cu'" diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index ef348e88d67..b58c6824ce6 100644 --- a/aiter/ops/mha.py +++ b/aiter/ops/mha.py @@ -16,7 +16,7 @@ is_experimental_enabled, ) from ..jit.utils.asm_guard import is_gfx1250_asm_supported, require_gfx1250_asm -from ..jit.utils.chip_info import get_cu_num, get_gfx +from ..jit.utils.chip_info import get_cu_num, get_gfx, get_gfx_runtime from ..jit.utils.mha_recipes import ( compose_mha_fwd_variant_suffix_and_filter, get_mha_varlen_prebuild_variants_by_names, @@ -690,7 +690,7 @@ def fmha_fwd_with_sink_varlen_asm( Contract details: * D64 / D128 take no strides: q/k/v are made dense, `out` must be dense. * D192x128 takes strides: q/k/v/out are used in place as views; the - dispatchers only send tensors `_fmha_varlen_asm_d192_operands_ok` accepts. + dispatchers only send tensors `_can_impl_fmha_varlen_asm_d192_strides` accepts. * `max_seqlen_q` is the maximum per-batch Q sequence length (caller- supplied, e.g. flash_attn_varlen convention) -- it sets the launch tile count; the kernel early-exits tiles beyond each batch's actual length. @@ -733,24 +733,272 @@ def fmha_fwd_with_sink_varlen_asm( return out, lse -def _fmha_varlen_asm_strided_ok(x: Tensor) -> bool: - """True if the D192x128 kernel can address `x` in place (mirrors the C++ check).""" - esz = x.element_size() - ts, hs = x.stride(0) * esz, x.stride(1) * esz - return x.stride(-1) == 1 and 0 < ts < (1 << 24) and 0 < hs <= 0x7FFFFFFF - - -def _fmha_varlen_asm_d192_operands_ok(q, k, v, out=None, dense_last_dim_later=False): - """True if the D192x128 kernel can use q/k/v (and `out`) in place. +def _can_impl_fmha_varlen_asm_d192_strides( + q, k, v, out=None, dense_last_dim_later=False +) -> bool: + """True if the D192x128 kernels can use q/k/v (and `out`) in place (mirrors the + C++ check): unit last-dim stride, token stride in (0, 2^24) bytes, head stride + in (0, 2^31) bytes. dense_last_dim_later: skip the last-dim check; `_flash_attn_varlen_forward` fixes it with `maybe_contiguous` and checks again. """ - ops = (q, k, v) if out is None else (q, k, v, out) - return all( - (dense_last_dim_later and x.stride(-1) != 1) or _fmha_varlen_asm_strided_ok(x) - for x in ops + for x in (q, k, v) if out is None else (q, k, v, out): + s = x.stride() + if s[-1] != 1: + if dense_last_dim_later: + continue + return False + esz = x.element_size() + if not (0 < s[0] * esz < (1 << 24) and 0 < s[1] * esz <= 0x7FFFFFFF): + return False + return True + + +# --------------------------------------------------------------------------- +# fmha_fwd_fp8_varlen_asm (gfx1250) — FP8 (e4m3fn) per-tensor-descale varlen +# forward, qk head dim 192 / v head dim 128, bf16 output. +# +# Separate from the bf16 varlen ASM path: its own kernarg ABI (136 B, the bf16 +# D192x128 strided layout with q/k/v descale pointers), its own C++ translation +# unit (csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu) and kernels +# (hsa/gfx1250/fmha_fwd_fp8_varlen). P is rounded to fp8 before PV and the row +# sum is taken over the rounded P, so the LSE is the log of that sum, an +# approximation of the exact log-sum-exp. +# --------------------------------------------------------------------------- +@compile_ops( + "module_fmha_fwd_fp8_varlen_asm", + fc_name="fmha_fwd_fp8_varlen_asm", + ffi_type="ctypes", +) +def _fmha_fwd_fp8_varlen_asm( + q: Tensor, + k: Tensor, + v: Tensor, + q_descale: Tensor, + k_descale: Tensor, + v_descale: Tensor, + out: Tensor, + lse: Tensor, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float, + is_causal: bool, + return_lse: bool, +) -> None: ... + + +def fmha_fwd_fp8_varlen_asm( + q: Tensor, + k: Tensor, + v: Tensor, + q_descale: Tensor, + k_descale: Tensor, + v_descale: Tensor, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + softmax_scale: float | None = None, + is_causal: bool = False, + return_lse: bool = False, + out: Tensor | None = None, +) -> tuple[Tensor, Tensor]: + """Public wrapper: varlen / packed FP8 ASM forward (gfx1250), D192x128. + + Layout is packed [token, head, dim] (THD), batch folded into the token axis: + * q : (total_q, nheads, 192) float8_e4m3fn + * k : (total_k, nheads_k, 192) float8_e4m3fn + * v : (total_k, nheads_k, 128) float8_e4m3fn + * q/k/v_descale : one fp32 element each (per-tensor dequant scale) + * out : (total_q, nheads, 128) bf16 + * lse : (total_q, nheads, 1) fp32, natural log + * cu_seqlens_q/k : int32 [batch+1] cumulative (cu[batch] == total) + + q/k/v/out are used in place through their strides (contiguous last dim; + `_can_impl_fmha_varlen_asm_d192_strides` lists the limits). `max_seqlen_q` sets + the launch tile count; with `max_seqlen_k`, the head count and the batch it + also picks the launch order (csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu). Causal is + bottom-right aligned. The kernel always writes `lse`. + """ + require_gfx1250_asm("fmha_fwd_fp8_varlen_asm") + if softmax_scale is None: + softmax_scale = q.size(-1) ** (-0.5) + cu_seqlens_q = cu_seqlens_q.to(torch.int32).contiguous() + cu_seqlens_k = cu_seqlens_k.to(torch.int32).contiguous() + + total_q, q_head_num, _qk_head_dim = q.shape + v_head_dim = v.size(2) + + if out is None: + out = torch.empty( + (total_q, q_head_num, v_head_dim), dtype=torch.bfloat16, device=q.device + ) + + lse = torch.empty((total_q, q_head_num, 1), dtype=torch.float32, device=q.device) + + _fmha_fwd_fp8_varlen_asm( + q, + k, + v, + q_descale, + k_descale, + v_descale, + out, + lse, + cu_seqlens_q, + cu_seqlens_k, + int(max_seqlen_q), + int(max_seqlen_k), + float(softmax_scale), + bool(is_causal), + bool(return_lse), ) + return out, lse + + +# --------------------------------------------------------------------------- +# gfx1250 varlen ASM forward: one gate for both kernels. +# +# flash_attn_varlen_func asks it to send a call to the ASM ahead of FlyDSL, +# _flash_attn_varlen_forward asks it to pick the kernel, and +# flash_attn_varlen_fp8_pertensor_func asks it to bypass Triton when CK is +# disabled. All of them pass the same arguments, so a call gets the same answer +# everywhere, and each site makes one call. Which accepted bf16 calls +# flash_attn_varlen_func still offers to FlyDSL first is a routing choice, kept +# apart in _prefer_flydsl_over_gfx1250_bf16_varlen_asm. +# --------------------------------------------------------------------------- +_GFX1250_VARLEN_ASM_BF16 = 1 # fmha_fwd_with_sink_varlen_asm +_GFX1250_VARLEN_ASM_FP8 = 2 # fmha_fwd_fp8_varlen_asm + + +def _probe_gfx1250_asm_device() -> bool: + # The running GPU, not the build target: get_gfx() follows GPU_ARCHS, which + # can name another arch (or several) than the device in use. + try: + arch = get_gfx_runtime() + except Exception: # noqa: BLE001 - no GPU to probe (e.g. a build host) + return False + return arch in ("gfx1250", "gfx1250-strict") and is_gfx1250_asm_supported() + + +# Probed once at import, like asm_guard._ASM_SUPPORTED: the dispatch sites read a +# module-level constant, which costs nothing per call and stays a plain constant +# under torch.compile. +_IS_GFX1250_ASM_DEVICE = _probe_gfx1250_asm_device() + + +def _select_gfx1250_fmha_fwd_varlen_asm( + q: Tensor, + k: Tensor, + v: Tensor, + out: Tensor | None, + *, + causal: bool, + max_seqlen_q: int, + max_seqlen_k: int, + window_size_left: int, + window_size_right: int, + sink_size: int, + sink_ptr: Tensor | None, + q_descale: Tensor | None, + k_descale: Tensor | None, + v_descale: Tensor | None, + dropout_p: float, + logits_soft_cap: float, + bias: Tensor | None, + alibi_slopes: Tensor | None, + block_table: Tensor | None, + cu_seqlens_q_padded: Tensor | None, + cu_seqlens_k_padded: Tensor | None, + return_softmax: bool, + dense_last_dim_later: bool = False, +) -> int: + """Which gfx1250 varlen ASM forward serves the call: _GFX1250_VARLEN_ASM_BF16, + _GFX1250_VARLEN_ASM_FP8, or 0 for neither. + + Windows and the sink-token span follow _flash_attn_varlen_forward and CK: a + negative value, or one at or beyond max_seqlen_k, is unlimited, and causal + ignores the right window. dense_last_dim_later: skip the last-dim stride + check, for callers ahead of the `maybe_contiguous` in + _flash_attn_varlen_forward. Cheap checks first, strides last. + """ + if not _IS_GFX1250_ASM_DEVICE: + return 0 + hdim_q, hdim_v = q.shape[-1], v.shape[-1] + if q.dtype == dtypes.bf16: + # fmha_fwd_with_sink_varlen_asm: D64 / D128 dense, D192x128 strided. + if q_descale is not None or k_descale is not None or v_descale is not None: + return 0 + if hdim_q == 192 and hdim_v == 128: + strided = True + elif hdim_q in (64, 128) and hdim_v == hdim_q: + strided = False + else: + return 0 + # Sink logits: the D64 binaries always read the sink buffer (a null one + # would be dereferenced); the D128 / D192x128 binaries never read it. + if (sink_ptr is not None) != (hdim_q == 64): + return 0 + kind = _GFX1250_VARLEN_ASM_BF16 + elif q.dtype == torch.float8_e4m3fn: + # fmha_fwd_fp8_varlen_asm: D192x128, per-tensor fp32 descales, bf16 out. + if k.dtype != q.dtype or v.dtype != q.dtype: + return 0 + if q.dim() != 3 or k.dim() != 3 or v.dim() != 3: + return 0 + if hdim_q != 192 or k.shape[-1] != 192 or hdim_v != 128: + return 0 + if sink_ptr is not None or (out is not None and out.dtype != torch.bfloat16): + return 0 + for d in (q_descale, k_descale, v_descale): + if ( + d is None + or d.dtype != torch.float32 + or d.numel() != 1 + or d.device != q.device + ): + return 0 + strided = True + kind = _GFX1250_VARLEN_ASM_FP8 + else: + return 0 + nhead_q, nhead_k = q.shape[-2], k.shape[-2] + if nhead_q <= 0 or nhead_k <= 0 or nhead_q % nhead_k != 0: + return 0 + # Attention options: plain or causal (bottom-right) attention only. + if max_seqlen_q <= 0 or return_softmax: + return 0 + if dropout_p != 0.0 or logits_soft_cap != 0.0: + return 0 + if 0 <= window_size_left < max_seqlen_k: + return 0 + if not causal and 0 <= window_size_right < max_seqlen_k: + return 0 + if sink_size != 0 and sink_size < max_seqlen_k: + return 0 + if bias is not None or alibi_slopes is not None or block_table is not None: + return 0 + if cu_seqlens_q_padded is not None or cu_seqlens_k_padded is not None: + return 0 + # Strides last: the most expensive check. + if strided and not _can_impl_fmha_varlen_asm_d192_strides( + q, k, v, out, dense_last_dim_later + ): + return 0 + return kind + + +def _prefer_flydsl_over_gfx1250_bf16_varlen_asm(q: Tensor, causal: bool) -> bool: + """bf16 calls the gfx1250 varlen ASM takes but flash_attn_varlen_func offers to + FlyDSL first: non-causal D64 / D128, and D64 / D128 / D192x128 with + AITER_ENABLE_EXPERIMENTAL (then the bshd and thd routers agree).""" + hdim_q = q.size(-1) + if is_experimental_enabled() and hdim_q in (64, 128, 192): + return True + return not causal and hdim_q in (64, 128) # --------------------------------------------------------------------------- @@ -3066,44 +3314,6 @@ def can_impl_fmha_v3_fwd(): ret = ret and not _fmha_kv_byte_extent_ge_u32(max_seqlen_k, k, v) return ret - def can_impl_fmha_fwd_with_sink_varlen_asm(): - # gfx1250 ASM bf16 packed/varlen forward (fmha_fwd_with_sink_varlen_asm). - # Packed THD (batch folded into the token axis); no dropout / swa / - # quant / alibi / bias / paged (block_table) / logits-soft-cap. Sink - # logits (per-Q-head fp32) supported; sink-token (sink_size) not. - ret = get_gfx() in ("gfx1250", "gfx1250-strict") and is_gfx1250_asm_supported() - ret = ret and (q.dtype == dtypes.bf16) - ret = ret and ( - (hdim_q in (64, 128) and hdim_v == hdim_q) - or (hdim_q == 192 and hdim_v == 128) - ) - # D192x128: strides the kernel cannot address go to another backend. - if hdim_q == 192 and hdim_v == 128: - ret = ret and _fmha_varlen_asm_d192_operands_ok(q, k, v, out) - ret = ret and (nhead_q % nhead_k == 0) - ret = ret and (not swa) - ret = ret and (sink_size == 0) - ret = ret and (alibi_slopes is None and bias is None) - ret = ret and (dropout_p == 0.0) - ret = ret and (logits_soft_cap == 0.0) - ret = ret and (block_table is None) - # The varlen ASM wrapper carries no physical-padding arrays; route any - # padded-cu request to CK (mha_varlen_fwd) which understands them. - ret = ret and (cu_seqlens_q_padded is None and cu_seqlens_k_padded is None) - ret = ret and (q_descale is None and k_descale is None and v_descale is None) - # Per-hdim sink eligibility (mirrors the fixed-batch path): - # D128 (`_rxy`) binaries compile ENABLE_SINK=0 and ignore the sink - # buffer, so routing a caller's sink_ptr to them would silently drop - # the sink term -- fall back to CK whenever sink_ptr is set. - # D64 (`_rxy_sink`) binaries compile ENABLE_SINK=1 and ALWAYS read - # SINK, so calling with sink_ptr=None would dereference a null pointer - # -- require an explicit sink for D64 and fall back to CK otherwise. - if hdim_q in (128, 192): - ret = ret and (sink_ptr is None) - elif hdim_q == 64: - ret = ret and (sink_ptr is not None) - return ret - def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen(): # OPUS gfx950 group/varlen D_QK=192 / D_V=128 bf16 forward. Enabled by DEFAULT # (no env). Packed THD q/k/v; supports KV padding (cu_seqlens_*_padded), @@ -3128,6 +3338,32 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen(): return ret q, k, v = [maybe_contiguous(x) for x in (q, k, v)] + gfx1250_asm = 0 + if _IS_GFX1250_ASM_DEVICE: + gfx1250_asm = _select_gfx1250_fmha_fwd_varlen_asm( + q, + k, + v, + out, + causal=causal, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + window_size_left=window_size_left, + window_size_right=window_size_right, + sink_size=sink_size, + sink_ptr=sink_ptr, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + dropout_p=dropout_p, + logits_soft_cap=logits_soft_cap, + bias=bias, + alibi_slopes=alibi_slopes, + block_table=block_table, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_k_padded=cu_seqlens_k_padded, + return_softmax=return_softmax, + ) if can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen(): # OPUS gfx950 group/varlen D=192 path. cu_seqlens_* are the REAL cumulative @@ -3155,13 +3391,12 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen(): ) S_dmask = torch.empty((0,), dtype=torch.float32, device=q.device) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) - elif can_impl_fmha_fwd_with_sink_varlen_asm(): + elif gfx1250_asm == _GFX1250_VARLEN_ASM_BF16: # gfx1250 packed/varlen ASM bf16 path (THD; D192x128 takes strided views). # softmax_scale is forwarded as-is (the kernel applies it internally to - # Q·K^T). sink_ptr is passed through verbatim; - # `can_impl_fmha_fwd_with_sink_varlen_asm` already enforces the per-hdim - # (D128 / D192x128 → no sink, D64 → sink) contract so we never feed a - # null sink to a D64 binary that unconditionally reads it. + # Q·K^T). sink_ptr is passed through verbatim; the gate already enforces + # the per-hdim (D128 / D192x128 → no sink, D64 → sink) contract so we never + # feed a null sink to a D64 binary that unconditionally reads it. out, lse_asm = fmha_fwd_with_sink_varlen_asm( q, k, @@ -3176,9 +3411,38 @@ def can_impl_fmha_fwd_hd192_v128_bf16_opus_varlen(): out, ) # The ASM kernel writes packed lse (total_q, nheads, 1); the varlen API - # convention (mha_varlen_fwd) is (nheads, total_q). Reshape so callers - # and the autograd backward see a consistent layout regardless of path. - softmax_lse = lse_asm.squeeze(-1).transpose(0, 1).contiguous() + # convention (mha_varlen_fwd) is (nheads, total_q). Made contiguous only + # when it is returned: FlashAttnVarlenFunc needs return_lse for backward. + softmax_lse = lse_asm.squeeze(-1).transpose(0, 1) + if return_lse: + softmax_lse = softmax_lse.contiguous() + S_dmask = torch.empty((0,), dtype=torch.float32, device=q.device) + rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) + elif gfx1250_asm == _GFX1250_VARLEN_ASM_FP8: + # gfx1250 packed/varlen ASM fp8 path (D192x128, per-tensor descales, + # strided views used in place). Output is bf16; softmax_scale is + # forwarded as-is and the descales are applied inside the kernel. + out, lse_asm = fmha_fwd_fp8_varlen_asm( + q, + k, + v, + q_descale, + k_descale, + v_descale, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + float(softmax_scale), + bool(causal), + bool(return_lse), + out, + ) + # packed (total_q, nheads, 1) -> the varlen API's (nheads, total_q), + # copied only when it is returned (as in the bf16 branch above) + softmax_lse = lse_asm.squeeze(-1).transpose(0, 1) + if return_lse: + softmax_lse = softmax_lse.contiguous() S_dmask = torch.empty((0,), dtype=torch.float32, device=q.device) rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) elif can_impl_fmha_v3_fwd(): @@ -3827,81 +4091,70 @@ def flash_attn_varlen_func( "FP8 attention requires q_descale, k_descale, and v_descale together" ) - # Try the PR3039 gfx1250 prefill ASM path before FlyDSL can claim it. - def can_try_gfx1250_fmha_fwd_with_sink_varlen_asm(): - # Keep this public-router gate intentionally narrow so the PR3039 - # prefill ASM path can be measured without changing decode or other - # FlyDSL/CK coverage. - if get_gfx() != "gfx1250" or not is_gfx1250_asm_supported(): - return False - if q.dtype != dtypes.bf16: - return False - hdim_q = q.shape[-1] - hdim_v = v.shape[-1] - nhead_q = q.shape[-2] - nhead_k = k.shape[-2] - is_hd192x128 = hdim_q == 192 and hdim_v == 128 - if not ((hdim_q in (64, 128) and hdim_v == hdim_q) or is_hd192x128): - return False - # D192x128: strides the kernel cannot address go to another backend. - if is_hd192x128 and not _fmha_varlen_asm_d192_operands_ok( - q, k, v, out, dense_last_dim_later=True - ): - return False - # Yield these to the FlyDSL m32x8 kernel below; without 192 the bshd and thd routers disagree. - if hdim_q in (64, 128, 192) and is_experimental_enabled(): - return False - if nhead_q % nhead_k != 0: - return False - if dropout_p != 0.0 or logits_soft_cap != 0.0: - return False - if not causal and not is_hd192x128: - return False - if window_size[0] != -1 or window_size[1] != -1: - return False - sink_size = window_size[2] if len(window_size) > 2 else 0 - if sink_size != 0: - return False - if bias is not None or alibi_slopes is not None or block_table is not None: - return False - if cu_seqlens_q_padded is not None or cu_seqlens_k_padded is not None: - return False - if hdim_q == 64: - return sink_ptr is not None - return sink_ptr is None - - if not has_all_descales and can_try_gfx1250_fmha_fwd_with_sink_varlen_asm(): - return FlashAttnVarlenFunc.apply( + # gfx1250 varlen ASM goes ahead of FlyDSL. _flash_attn_varlen_forward asks + # the same gate with the arguments FlashAttnVarlenFunc passes it, so an + # accepted call is served by that ASM kernel. + if _IS_GFX1250_ASM_DEVICE: + gfx1250_asm = _select_gfx1250_fmha_fwd_varlen_asm( q, k, v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - min_seqlen_q, - dropout_p, - softmax_scale, - logits_soft_cap, - causal, - window_size, - bias, - alibi_slopes, - deterministic, - return_lse, - return_attn_probs, - block_table, out, - torch.is_grad_enabled(), - cu_seqlens_q_padded, - cu_seqlens_k_padded, - True, - how_v3_bf16_cvt, - sink_ptr, - q_descale, - k_descale, - v_descale, + causal=causal, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + window_size_left=window_size[0], + window_size_right=window_size[1], + sink_size=window_size[2] if len(window_size) > 2 else 0, + sink_ptr=sink_ptr, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + dropout_p=dropout_p, + logits_soft_cap=logits_soft_cap, + bias=bias, + alibi_slopes=alibi_slopes, + block_table=block_table, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_k_padded=cu_seqlens_k_padded, + return_softmax=return_attn_probs and dropout_p > 0, + dense_last_dim_later=True, # _flash_attn_varlen_forward fixes it ) + if gfx1250_asm == _GFX1250_VARLEN_ASM_FP8 or ( + gfx1250_asm == _GFX1250_VARLEN_ASM_BF16 + and not _prefer_flydsl_over_gfx1250_bf16_varlen_asm(q, causal) + ): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + min_seqlen_q, + dropout_p, + softmax_scale, + logits_soft_cap, + causal, + window_size, + bias, + alibi_slopes, + deterministic, + return_lse, + return_attn_probs, + block_table, + out, + torch.is_grad_enabled(), + cu_seqlens_q_padded, + cu_seqlens_k_padded, + True, + how_v3_bf16_cvt, + sink_ptr, + q_descale, + k_descale, + v_descale, + ) if ( _FLYDSL_AVAILABLE @@ -4375,7 +4628,38 @@ def flash_attn_varlen_fp8_pertensor_func( softmax_scale=None, sink_ptr=None, ): - if not ENABLE_CK and sink_ptr is None: + # The gfx1250 FP8 D192x128 ASM kernel is reached through + # _flash_attn_varlen_forward, which takes it before CK; with CK disabled, + # go there only for calls that kernel accepts (same gate, same arguments). + use_fp8_varlen_asm = _IS_GFX1250_ASM_DEVICE and ( + _select_gfx1250_fmha_fwd_varlen_asm( + q, + k, + v, + None, + causal=causal, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + window_size_left=int(window_size[0]), + window_size_right=int(window_size[1]), + sink_size=int(window_size[2]) if len(window_size) == 3 else 0, + sink_ptr=sink_ptr, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + dropout_p=0.0, + logits_soft_cap=logits_soft_cap, + bias=None, + alibi_slopes=None, + block_table=None, + cu_seqlens_q_padded=None, + cu_seqlens_k_padded=None, + return_softmax=False, + dense_last_dim_later=True, # _flash_attn_varlen_forward fixes it + ) + == _GFX1250_VARLEN_ASM_FP8 + ) + if not ENABLE_CK and sink_ptr is None and not use_fp8_varlen_asm: from .triton.attention.mha_v3 import ( flash_attn_varlen_func as flash_attn_varlen_func_v3_triton, ) diff --git a/csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu b/csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu new file mode 100644 index 00000000000..4970ec06462 --- /dev/null +++ b/csrc/py_itfs_cu/asm_fmha_fwd_fp8_varlen.cu @@ -0,0 +1,344 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +// +// ASM FMHA forward, VARLEN / packed, FP8 (e4m3fn) with per-tensor descales, gfx1250. +// qk head dim 192, v head dim 128; output bf16. +// +// q : (total_q, nheads, 192) fp8 e4m3fn +// k : (total_k, nheads_k, 192) fp8 e4m3fn +// v : (total_k, nheads_k, 128) fp8 e4m3fn +// out : (total_q, nheads, 128) bf16 +// lse : packed [total_q, nheads] fp32, natural log (caller may shape (total_q, nheads, 1)) +// q/k/v_descale : fp32, one element each (per-tensor), read by the kernel from device memory +// cu_seqlens_q/k : int32 [batch+1] cumulative, cu[batch] == total +// +// q/k/v/out are addressed through their token / head strides (contiguous last dim), so views +// of a larger buffer work in place. Causal masking is bottom-right aligned. +// +// The softmax probabilities P are rounded to fp8 before the PV product, and the row sum (the +// denominator of O and the LSE) is the sum of those rounded P: O is consistent with the P actually +// used, and the LSE is the log of that same sum, an approximation of the exact log-sum-exp. +// +// Memory-allocation policy: all tensors are allocated by the Python caller. This entry point +// only fills the kernel arguments and launches -- no GPU allocation, no torch dependency. +#include "aiter_tensor.h" +#include "aiter_ctypes_error.h" +#include "aiter_hip_common.h" // HipDeviceGuard, AiterAsmKernel, ... +#include "asm_fmha_fwd_fp8_varlen_configs.hpp" +#include +#include +#include +#include +#include + +// Kernel argument block (136 B). The first 120 B mirror the bf16 D192x128 strided varlen ABI +// (FmhaFwdVarlenStrideKernelArgs) with the sink pointer slot at 0x50 holding the q descale +// pointer; k/v descale pointers follow at 0x78 / 0x80. The kernel preloads the first 120 B +// (s2..s31) and loads the last 16 B itself. +#pragma pack(push, 1) +struct FmhaFwdFp8VarlenKernelArgs +{ + void* d_addr; // off 0x00 O output (total_q, nheads, 128) bf16 + const void* q_addr; // off 0x08 Q (total_q, nheads, 192) + const void* k_addr; // off 0x10 K (total_k, nheads_k, 192) + const void* v_addr; // off 0x18 V (total_k, nheads_k, 128) + void* lse_addr; // off 0x20 LSE packed [total_q, nheads] + const void* qseq_addr; // off 0x28 cu_seqlens_q int32[batch+1] + const void* kseq_addr; // off 0x30 cu_seqlens_k int32[batch+1] + float scalar; // off 0x38 softmax_scale (descales are applied by the kernel) + int gqa; // off 0x3C nheads / nheads_k + int q_head_num; // off 0x40 nheads + int opt; // off 0x44 bit1 double_q | bit2 remap_xy (bit0 ignored) + int lse; // off 0x48 1 = write LSE + int max_q_len; // off 0x4C max over batches of q seqlen (dispatch basis) + const void* q_descale_addr; // off 0x50 fp32 [1] + int q_hs; // off 0x58 Q head stride (bytes) + int k_hs; // off 0x5C K head stride (bytes) + int v_hs; // off 0x60 V head stride (bytes) + int d_hs; // off 0x64 out head stride (bytes) + int q_seqs; // off 0x68 Q token stride (bytes) + int k_seqs; // off 0x6C K token stride (bytes) + int v_seqs; // off 0x70 V token stride (bytes) + int d_seqs; // off 0x74 out token stride (bytes) + const void* k_descale_addr; // off 0x78 fp32 [1] + const void* v_descale_addr; // off 0x80 fp32 [1] +}; +#pragma pack(pop) +static_assert(sizeof(FmhaFwdFp8VarlenKernelArgs) == 0x88, + "fmha_fwd_fp8_varlen_asm: FmhaFwdFp8VarlenKernelArgs must be 136B packed"); + +// ---- helpers --------------------------------------------------------------- + +// Kernel selection: (dtype, hdim_q, hdim_v, mask); mask = is_causal (the causal kernel is a +// separate build). +static std::string get_heuristic_kernel_fmha_fwd_fp8_varlen(const std::string& dtype, + int hdim_q, + int hdim_v, + int mask_flag, + const std::string& arch_id, + CFG* cfgs) +{ + for (const auto& el : *cfgs) + { + if (el.first.find(arch_id) != 0) continue; + const auto& cfg = el.second; + if (cfg.dtype != dtype) continue; + if (cfg.hdim_q != hdim_q) continue; + if (cfg.hdim_v != hdim_v) continue; + if (cfg.mask != mask_flag) continue; + return el.first; + } + AITER_CHECK(false, + "fmha_fwd_fp8_varlen_asm: no kernel for dtype=", dtype, + " hdim_q=", hdim_q, " hdim_v=", hdim_v, + " mask=", mask_flag, + " arch=", arch_id); + return ""; +} + +// The kernel steps the KV loop by 64 * token stride as a 32-bit value and keeps the strides in +// 32-bit registers; token stride < 2^24 B keeps the same margin as the bf16 D192x128 kernel. +static void check_strided_thd(const aiter_tensor_t* t, const char* name) +{ + const int64_t esz = (int64_t)t->element_size(); + const int64_t ts = t->stride(0) * esz; + const int64_t hs = t->stride(1) * esz; + AITER_CHECK(t->stride(-1) == 1, + "fmha_fwd_fp8_varlen_asm: ", name, " must have contiguous last dim"); + AITER_CHECK(ts > 0 && ts < (int64_t(1) << 24) && hs > 0 && hs <= INT32_MAX, + "fmha_fwd_fp8_varlen_asm: ", name, " strides out of range (token ", ts, + " B must be in (0, 2^24), head ", hs, " B must be in (0, INT32_MAX])"); +} + +static void check_descale(const aiter_tensor_t* t, const char* name, int device_id) +{ + AITER_CHECK(t->dtype() == AITER_DTYPE_fp32 && t->numel() == 1, + "fmha_fwd_fp8_varlen_asm: ", name, " must be a single fp32 value (per-tensor)"); + AITER_CHECK(t->device_id == device_id, + "fmha_fwd_fp8_varlen_asm: ", name, " must be on the same GPU as q"); +} + +static int cu_count(int device) +{ + static std::atomic cache[64]; + if (device < 0 || device >= 64) + device = 0; + int n = cache[device].load(std::memory_order_relaxed); + if (n <= 0) + { + if (hipDeviceGetAttribute(&n, hipDeviceAttributeMultiprocessorCount, device) != hipSuccess || n <= 0) + n = 256; + cache[device].store(n, std::memory_order_relaxed); + } + return n; +} + +// Fraction of the last round of workgroups that is busy when n workgroups run s at a time. +static double wave_fill(long n, int s) +{ + if (n <= 0) + return 1.0; + const long rounds = (n + s - 1) / s; + return (double)n / ((double)rounds * s); +} + +// Launch options (bit1 double_q, bit2 remap_xy). One 4-wave workgroup runs per CU. +// double_q: a workgroup runs Q tiles g and tiles-1-g one after the other, so the grid has half as +// many workgroups. Causal: every workgroup gets the same work (a heavy tile plus a light one). +// Non-causal: the two tiles share K/V in L2 -- a gain unless halving the grid leaves the last +// round of workgroups emptier, so the fill of the last round is compared with and without it. +// Causal without double_q relies on the heaviest-first launch order instead (remap_xy, below). +// remap_xy: x indexes heads, so workgroups launched together are different heads of one Q tile. +// It helps while the heads' K/V together are small (q heads x kv length under the threshold +// below); with many heads over long sequences they stream K/V through L2 at once and it hurts. +// Causal without double_q always uses it: it launches the heavy tiles of every head first. +// The thresholds come from timing the four options over head counts and sequence lengths on gfx1250, +// weighted towards many heads with long sequences. +// AITER_FMHA_FP8_VARLEN_ASM_OPT= overrides the choice (bits 1 and 2 only). +static int select_opt(int batch, int max_seqlen_q, int max_seqlen_k, int q_head_num, int is_causal, int device) +{ + const char* env = std::getenv("AITER_FMHA_FP8_VARLEN_ASM_OPT"); + if (env && *env) + return (int)std::strtol(env, nullptr, 0) & 0x6; + const int s = cu_count(device); + const long tiles = (max_seqlen_q + 255) / 256; + const long wg = tiles * q_head_num * batch; // workgroups without double_q + const long wg_dq = (tiles + 1) / 2 * q_head_num * batch; // with double_q + const bool remap = (long)q_head_num * max_seqlen_k <= (1L << 18); + const bool dq = is_causal ? wave_fill(wg_dq, s) >= 0.93 + : wave_fill(wg_dq, s) * 1.03 >= wave_fill(wg, s); + if (dq) + return 0x2 | (remap ? 0x4 : 0x0); + return is_causal ? 0x4 : 0x0; +} + +// ---- main entry ------------------------------------------------------------ + +AITER_CTYPES_ERROR_DEF + +AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( + fmha_fwd_fp8_varlen_asm, + (aiter_tensor_t* q, + aiter_tensor_t* k, + aiter_tensor_t* v, + aiter_tensor_t* q_descale, + aiter_tensor_t* k_descale, + aiter_tensor_t* v_descale, + aiter_tensor_t* out, + aiter_tensor_t* lse, + aiter_tensor_t* cu_seqlens_q, + aiter_tensor_t* cu_seqlens_k, + int max_seqlen_q, + int max_seqlen_k, + float softmax_scale, + int is_causal, + int return_lse, + hipStream_t stream), + (q, k, v, q_descale, k_descale, v_descale, out, lse, cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_k, softmax_scale, is_causal, return_lse, stream)) +{ + // ---- null safety --------------------------------------------------------- + AITER_CHECK(q && k && v && out && lse && cu_seqlens_q && cu_seqlens_k, + "fmha_fwd_fp8_varlen_asm: q/k/v/out/lse/cu_seqlens_q/cu_seqlens_k must all be non-null"); + AITER_CHECK(q_descale && k_descale && v_descale, + "fmha_fwd_fp8_varlen_asm: q/k/v_descale must all be non-null"); + + HipDeviceGuard device_guard{q->device_id}; + + // ---- arch + dtype validation ---------------------------------------------- + const std::string arch_id = get_gpu_arch(); + AITER_CHECK((arch_id == "gfx1250" || arch_id == "gfx1250-strict"), + "fmha_fwd_fp8_varlen_asm: only supported on gfx1250, got ", arch_id); + + AITER_CHECK(q->dtype() == AITER_DTYPE_fp8 && + k->dtype() == AITER_DTYPE_fp8 && + v->dtype() == AITER_DTYPE_fp8, + "fmha_fwd_fp8_varlen_asm: q/k/v must be fp8 (e4m3fn)"); + AITER_CHECK(out->dtype() == AITER_DTYPE_bf16, + "fmha_fwd_fp8_varlen_asm: out must be bf16"); + AITER_CHECK(lse->dtype() == AITER_DTYPE_fp32, + "fmha_fwd_fp8_varlen_asm: lse must be fp32"); + AITER_CHECK(cu_seqlens_q->dtype() == AITER_DTYPE_i32 && + cu_seqlens_k->dtype() == AITER_DTYPE_i32, + "fmha_fwd_fp8_varlen_asm: cu_seqlens_q/k must be int32"); + check_descale(q_descale, "q_descale", q->device_id); + check_descale(k_descale, "k_descale", q->device_id); + check_descale(v_descale, "v_descale", q->device_id); + + // ---- shape extraction (packed thd) ---------------------------------------- + AITER_CHECK(q->dim() == 3 && k->dim() == 3 && v->dim() == 3, + "fmha_fwd_fp8_varlen_asm: q/k/v must be 3-D packed tensors (total, head, dim)"); + + const int total_q = (int)q->size(0); + const int q_head_num = (int)q->size(1); + const int qk_head_dim = (int)q->size(2); + + const int total_k = (int)k->size(0); + const int kv_head_num = (int)k->size(1); + const int v_head_dim = (int)v->size(2); + + AITER_CHECK(qk_head_dim == 192 && v_head_dim == 128 && (int)k->size(2) == qk_head_dim, + "fmha_fwd_fp8_varlen_asm: only (qk_head_dim, v_head_dim) = (192, 128) is supported; got (", + qk_head_dim, ",", v_head_dim, ")"); + AITER_CHECK((int)v->size(0) == total_k, "fmha_fwd_fp8_varlen_asm: v total_k mismatch with k"); + AITER_CHECK((int)v->size(1) == kv_head_num, "fmha_fwd_fp8_varlen_asm: v head_num mismatch with k"); + AITER_CHECK(q_head_num > 0 && kv_head_num > 0, + "fmha_fwd_fp8_varlen_asm: q/kv head counts must be positive"); + AITER_CHECK(q_head_num % kv_head_num == 0, + "fmha_fwd_fp8_varlen_asm: q_head_num must be a multiple of kv_head_num"); + + AITER_CHECK(out->dim() == 3 && + (int)out->size(0) == total_q && (int)out->size(1) == q_head_num && + (int)out->size(2) == v_head_dim, + "fmha_fwd_fp8_varlen_asm: out shape must be [total_q, q_head_num, 128]"); + AITER_CHECK(lse->dim() >= 2 && + (int)lse->size(0) == total_q && (int)lse->size(1) == q_head_num, + "fmha_fwd_fp8_varlen_asm: lse leading dims must be [total_q, q_head_num]"); + AITER_CHECK(lse->is_contiguous(), + "fmha_fwd_fp8_varlen_asm: lse must be contiguous"); + + check_strided_thd(q, "q"); + check_strided_thd(k, "k"); + check_strided_thd(v, "v"); + check_strided_thd(out, "out"); + + AITER_CHECK(cu_seqlens_q->dim() == 1 && cu_seqlens_k->dim() == 1, + "fmha_fwd_fp8_varlen_asm: cu_seqlens_q/k must be 1-D"); + const int batch = (int)cu_seqlens_q->size(0) - 1; + AITER_CHECK(batch >= 1, + "fmha_fwd_fp8_varlen_asm: cu_seqlens_q must have length batch+1 (>=2)"); + AITER_CHECK((int)cu_seqlens_k->size(0) == batch + 1, + "fmha_fwd_fp8_varlen_asm: cu_seqlens_k length must match cu_seqlens_q"); + AITER_CHECK(max_seqlen_q > 0, + "fmha_fwd_fp8_varlen_asm: max_seqlen_q must be > 0"); + + const int mask_flag = is_causal ? 1 : 0; + const int opt = select_opt(batch, max_seqlen_q, max_seqlen_k, q_head_num, is_causal, q->device_id); + + // ---- kernel args ------------------------------------------------------------- + FmhaFwdFp8VarlenKernelArgs args; + memset(&args, 0, sizeof(args)); + args.d_addr = out->data_ptr(); + args.q_addr = q->data_ptr(); + args.k_addr = k->data_ptr(); + args.v_addr = v->data_ptr(); + args.lse_addr = lse->data_ptr(); + args.qseq_addr = cu_seqlens_q->data_ptr(); + args.kseq_addr = cu_seqlens_k->data_ptr(); + args.scalar = softmax_scale; + args.gqa = q_head_num / kv_head_num; + args.q_head_num = q_head_num; + args.opt = opt; + args.lse = return_lse ? 1 : 0; + args.max_q_len = max_seqlen_q; + args.q_descale_addr = q_descale->data_ptr(); + args.q_hs = (int)(q->stride(1) * q->element_size()); + args.k_hs = (int)(k->stride(1) * k->element_size()); + args.v_hs = (int)(v->stride(1) * v->element_size()); + args.d_hs = (int)(out->stride(1) * out->element_size()); + args.q_seqs = (int)(q->stride(0) * q->element_size()); + args.k_seqs = (int)(k->stride(0) * k->element_size()); + args.v_seqs = (int)(v->stride(0) * v->element_size()); + args.d_seqs = (int)(out->stride(0) * out->element_size()); + args.k_descale_addr = k_descale->data_ptr(); + args.v_descale_addr = v_descale->data_ptr(); + size_t arg_size = sizeof(args); + + // ---- kernel selection ---------------------------------------------------------- + const std::string dtype = "fp8"; + CFG* cfg_map = &cfg_fmha_fwd_fp8_varlen; + static SynchronizedCache impl_ptr_map; + + const std::string kernel_key = get_heuristic_kernel_fmha_fwd_fp8_varlen( + dtype, qk_head_dim, v_head_dim, mask_flag, arch_id, cfg_map); + auto it = cfg_map->find(kernel_key); + AITER_CHECK(it != cfg_map->end(), + "fmha_fwd_fp8_varlen_asm: kernel not found in CFG: ", kernel_key); + + const char* name = it->second.knl_name.c_str(); + const char* co_name = it->second.co_name.c_str(); + AiterAsmKernel* impl_ptr = &impl_ptr_map.get_or_create( + name, [&]() { return AiterAsmKernel(name, co_name); }); + + // ---- launch ------------------------------------------------------------------------ + // x = Q tiles of 256 rows over max_q_len (halved by double_q), y = q heads, z = batch; + // remap_xy swaps x and y. 4 waves of 32 lanes per workgroup. Tiles beyond a batch's own + // length exit early. + const int sub_Q = 256; + const int bdx = 128; + const int q_tile_count = (max_seqlen_q + sub_Q - 1) / sub_Q; + const int gdx = (opt & 0x2) ? (q_tile_count + 1) / 2 : q_tile_count; + const int gdy = q_head_num; + const int gdz = batch; + const bool remap = (opt & 0x4) != 0; + + impl_ptr->launch_kernel({&args, + &arg_size, + remap ? gdy : gdx, + remap ? gdx : gdy, + gdz, + bdx, + 1, + 1, + stream}); +} diff --git a/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_mask_varlen.co b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_mask_varlen.co new file mode 100755 index 00000000000..fc31cdc8586 Binary files /dev/null and b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_mask_varlen.co differ diff --git a/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_varlen.co b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_varlen.co new file mode 100755 index 00000000000..5071d0f4e82 Binary files /dev/null and b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fp8_pertensor_hd192x128_256x64_varlen.co differ diff --git a/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fwd_fp8_varlen.csv b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fwd_fp8_varlen.csv new file mode 100644 index 00000000000..6bbabd4650c --- /dev/null +++ b/hsa/gfx1250/fmha_fwd_fp8_varlen/fmha_fwd_fp8_varlen.csv @@ -0,0 +1,3 @@ +dtype,hdim_q,hdim_v,mask,knl_name,co_name +fp8,192,128,0,_ZN5aiter42fmha_fp8_pertensor_hd192x128_256x64_varlenE,fmha_fp8_pertensor_hd192x128_256x64_varlen.co +fp8,192,128,1,_ZN5aiter47fmha_fp8_pertensor_hd192x128_256x64_mask_varlenE,fmha_fp8_pertensor_hd192x128_256x64_mask_varlen.co diff --git a/op_tests/test_fmha_fwd_with_sink_varlen_asm.py b/op_tests/test_fmha_fwd_with_sink_varlen_asm.py index f4578999649..202c915fd7d 100644 --- a/op_tests/test_fmha_fwd_with_sink_varlen_asm.py +++ b/op_tests/test_fmha_fwd_with_sink_varlen_asm.py @@ -1,9 +1,15 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness + perf tests for fmha_fwd_with_sink_varlen_asm (BF16 ASM, gfx1250). +"""Correctness + perf tests for the gfx1250 varlen ASM forward kernels: +fmha_fwd_with_sink_varlen_asm (BF16) and fmha_fwd_fp8_varlen_asm (FP8 D192x128). -Public API: aiter.flash_attn_varlen_func (the path the model calls) -Ops layer: aiter.fmha_fwd_with_sink_varlen_asm (low-level, packed/varlen) +BF16 public API: aiter.flash_attn_varlen_func (the path the model calls) + ops layer: aiter.fmha_fwd_with_sink_varlen_asm (low-level, packed/varlen) +FP8 public API: aiter.flash_attn_varlen_fp8_pertensor_func, + aiter.flash_attn_varlen_func(..., q/k/v_descale=...) + ops layer: aiter.fmha_fwd_fp8_varlen_asm (per-tensor fp32 descales, + qk 192 / v 128, bf16 output; its tables and tolerances are in + the FP8 section below; --dtype picks bf16 / fp8) Built to the aiter op-test standard (see .claude/skills/aiter-op-test): mirror test_quant.py — @benchmark + run_perftest candidate loop, a torch reference, @@ -30,7 +36,8 @@ constraint (sub_K=128 with the border mask compiled in). Strided inputs: D192x128 takes q/k/v/out views (e.g. v = kv[..., 128:]); they are -checked against the reference and bit-for-bit against contiguous inputs. +checked against the reference and bit-for-bit against contiguous inputs (BF16 and +FP8 alike). """ import argparse @@ -42,7 +49,7 @@ import torch import aiter -from aiter import dtypes +from aiter import dtypes, per_tensor_quant from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.test_common import benchmark, checkAllclose, run_perftest @@ -318,11 +325,20 @@ def test_fmha_fwd_with_sink_varlen_asm( return ret +def _nan_buffer(shape, dtype): + """NaN-filled buffer, so a read outside the view shows up (fp8: 0x7f, the + e4m3fn NaN).""" + if dtype == dtypes.fp8: + return torch.full(shape, 0x7F, dtype=torch.uint8).view(dtype) + return torch.full(shape, float("nan"), dtype=dtype) + + def _as_view(x, head_pad=0, token_pad=0, base_off=0): - """Strided copy of `x` in a NaN buffer; odd pads/offset give 2-byte-aligned rows.""" + """Strided copy of `x` in a NaN buffer; odd pads/offset give rows aligned + only to the element size (2 B for bf16, 1 B for fp8).""" t, h, d = x.shape row = h * (d + head_pad) + token_pad - buf = torch.full((base_off + t * row,), float("nan"), dtype=x.dtype) + buf = _nan_buffer((base_off + t * row,), x.dtype) view = buf.as_strided((t, h, d), (row, d + head_pad, 1), base_off) view.copy_(x) return view @@ -334,12 +350,13 @@ def _strided_inputs(layout, q, k, v): hk, dv = k.size(1), v.size(2) if layout == "kv_split": # DeepSeek-style: v is the second half of a [.., k_nope | v] projection. - kv = torch.full((t, hk, dv + dv), float("nan"), dtype=v.dtype) + kv = _nan_buffer((v.size(0), hk, dv + dv), v.dtype) kv[..., dv:] = v return q, k, kv[..., dv:] if layout == "fused_qkv": # One row per token holds every q, k and v head back to back. - row = torch.full((t, hq * dq + hk * dq + hk * dv), float("nan"), dtype=q.dtype) + assert k.size(0) == t, "fused_qkv needs equal q and k token counts" + row = _nan_buffer((t, hq * dq + hk * dq + hk * dv), q.dtype) qv = row[:, : hq * dq].view(t, hq, dq) kvv = row[:, hq * dq : hq * dq + hk * dq].view(t, hk, dq) vv = row[:, hq * dq + hk * dq :].view(t, hk, dv) @@ -354,7 +371,7 @@ def _strided_inputs(layout, q, k, v): _as_view(v, head_pad=128), ) if layout == "odd": - # Odd-element strides and offsets: rows start only 2-byte aligned. + # Odd-element strides and offsets: rows start only element-size aligned. return ( _as_view(q, head_pad=1, token_pad=3), _as_view(k, head_pad=1, token_pad=1), @@ -531,6 +548,93 @@ def spy(*args, **kwargs): } +# Attention options _select_gfx1250_fmha_fwd_varlen_asm checks the same way for +# both kernels. The gate tests use one sequence of 2 query and 2 key tokens, so +# max_seqlen_k is 2. name -> (gate arguments, accepted) +_ASM_GATE_OPTION_CASES = { + "causal": ({"causal": True}, True), + "causal, right window 0": ({"causal": True, "window_size_right": 0}, True), + "left window >= max_seqlen_k": ({"window_size_left": 2}, True), + "non-causal, right window 0": ({"window_size_right": 0}, False), + "sliding window": ({"window_size_left": 1}, False), + "sink tokens": ({"sink_size": 1}, False), + "dropout": ({"dropout_p": 0.1}, False), + "soft-cap": ({"logits_soft_cap": 30.0}, False), + "no query rows": ({"max_seqlen_q": 0}, False), +} + + +def _asm_gate_args(case): + """Attention-option arguments of the gate for `case`: non-causal plain + attention unless the case says otherwise.""" + args = { + "causal": False, + "max_seqlen_q": 2, + "max_seqlen_k": 2, + "window_size_left": -1, + "window_size_right": -1, + "sink_size": 0, + "dropout_p": 0.0, + "logits_soft_cap": 0.0, + "bias": None, + "alibi_slopes": None, + "block_table": None, + "cu_seqlens_q_padded": None, + "cu_seqlens_k_padded": None, + "return_softmax": False, + } + args.update(_ASM_GATE_OPTION_CASES.get(case, ({}, None))[0]) + return args + + +# name -> served by the bf16 kernel (_select_gfx1250_fmha_fwd_varlen_asm) +_BF16_GATE_CASES = { + "D192x128": True, + "D128": True, + "D64 with sink logits": True, + "D64 without sink logits": False, + "D128 with sink logits": False, + "D192x128 with sink logits": False, + "head dim 256": False, + "fp32 descales": False, + "zero kv heads": False, + **{c: ok for c, (_, ok) in _ASM_GATE_OPTION_CASES.items()}, +} + + +@benchmark() +def test_fmha_fwd_with_sink_varlen_asm_gate(case): + """The gate every dispatch site asks must pick the bf16 kernel for exactly what + it serves. Option cases use D192x128. Not timed: nothing is launched.""" + if case.startswith("D64"): + d, dv = 64, 64 + elif case.startswith("D128"): + d, dv = 128, 128 + elif case == "head dim 256": + d, dv = 256, 256 + else: + d, dv = 192, 128 + q, k, v, _ = make_varlen_packed([2], 4, 4, d, dv) + if case == "zero kv heads": + k, v = k[:, :0], v[:, :0] + sink = torch.zeros(4, dtype=dtypes.fp32) if "with sink" in case else None + descale = torch.ones(1, dtype=dtypes.fp32) if case == "fp32 descales" else None + mha = aiter.ops.mha + routed = mha._GFX1250_VARLEN_ASM_BF16 == mha._select_gfx1250_fmha_fwd_varlen_asm( + q, + k, + v, + None, + sink_ptr=sink, + q_descale=descale, + k_descale=descale, + v_descale=descale, + **_asm_gate_args(case), + ) + assert routed == _BF16_GATE_CASES[case], f"{case}: gate returned {routed}" + return {"gfx": get_gfx(), "routed to asm": routed} + + _STRIDED_LAYOUTS = [ "kv_split", "fused_qkv", @@ -584,6 +688,708 @@ def test_fmha_fwd_with_sink_varlen_asm_perf( return ret +# --------------------------------------------------------------------------- +# FP8 D192x128 (fmha_fwd_fp8_varlen_asm): fp8 e4m3fn q/k/v with per-tensor +# fp32 descales, bf16 output. +# +# Candidates (all must reach the ASM kernel and agree bit for bit): +# public -> aiter.flash_attn_varlen_fp8_pertensor_func +# varlen -> aiter.flash_attn_varlen_func(..., q/k/v_descale=...) (with LSE) +# ops -> aiter.fmha_fwd_fp8_varlen_asm +# +# Reference: fp32 on the dequantized fp8 inputs, emulating the kernel's online +# softmax block by block (as the poc host's pemu=2): keys in blocks of the +# kernel's KV tile, query rows in groups of one wave; the first block, or a block +# whose max exceeds the running max by more than FP8_TAU (exp2 units) in any row +# of the group, moves the whole group's max and rescales; P = exp(S - running max) +# is rounded to e4m3 and the row sum is the sum of the rounded P (so is the LSE). +# What is left between kernel and reference is the summation order and the +# reduced precision of S inside the fp8 WMMA, which now and then puts a P on the +# other side of an e4m3 rounding boundary. +# * O: an element is "bad" past FP8_ATOL + FP8_RTOL * |ref|. One such flip can +# shift a whole row with few visible keys, so the check counts rows (a token +# of a head) holding bad elements and allows a few per thousand; a layout / +# mask / rescale bug breaks whole 16-row tiles or every head. +# * cos >= FP8_COS_MIN over the whole output. +# * LSE: within FP8_LSE_EMU_TOL of the emulated one, for all but as many rows as +# O may have bad ones (a flipped P moves the sum of a short row); rows that see +# no key are -inf. The distance to the exact log-sum-exp is only reported: +# it is the rounding of P, which has no fixed bound -- P far below the row's +# max fall into e4m3's subnormal range or round to zero, and a row whose +# weight sits mostly in such P keeps a correspondingly smaller sum. +# --------------------------------------------------------------------------- + +try: # FlyDSL's gfx1250 fp8 kernel, a reference point in the fp8 perf table + from aiter.ops.flydsl.fmha_kernels import flydsl_flash_attn_varlen_fwd +except Exception: # noqa: BLE001 -- optional; the table just goes without it + flydsl_flash_attn_varlen_fwd = None + +FP8_ATOL, FP8_RTOL = 4e-2, 1e-2 +FP8_COS_MIN = 0.9999 +FP8_COS_MIN_EXACT = 0.999 # against the exact softmax: the rounding of P alone moves O +FP8_LSE_EMU_TOL = 1e-2 +FP8_TAU = 8.0 # the kernel's lazy-rescale threshold (log2 units) +_FP8_VIAS = ["public", "varlen", "ops"] + + +def _cu(seqlens): + return torch.tensor( + [0] + torch.tensor(seqlens).cumsum(0).tolist(), dtype=dtypes.i32 + ) + + +def _quant(x): + """Per-tensor fp8 quantization; an empty tensor (no keys at all) gets descale 1.""" + if x.numel() == 0: + return x.to(dtypes.fp8), torch.ones(1, dtype=dtypes.fp32) + return per_tensor_quant(x, quant_dtype=dtypes.fp8) + + +def make_fp8_varlen(seqlens_q, seqlens_k, hq, hk, seed=0): + """Packed fp8 q/k/v quantized from randn with per-tensor descales.""" + torch.manual_seed(seed) + tq, tk = sum(seqlens_q), sum(seqlens_k) + q8, qd = _quant(torch.randn(tq, hq, 192, dtype=dtypes.bf16)) + k8, kd = _quant(torch.randn(tk, hk, 192, dtype=dtypes.bf16)) + v8, vd = _quant(torch.randn(tk, hk, 128, dtype=dtypes.bf16)) + return q8, k8, v8, qd, kd, vd, _cu(seqlens_q), _cu(seqlens_k) + + +def _emulate_block(s, scale_log2e, causal_vis, vf): + """Kernel-order online softmax for one sequence: s (h, sq, sk) natural-log + scores with -inf masked, vf (h, sk, d) values. Returns O (sq, h, d), LSE of + the rounded P (h, sq).""" + h, sq, sk = s.shape + g = (sq + 63) // 64 + pad = g * 64 - sq + s = torch.nn.functional.pad(s, (0, 0, 0, pad), value=float("-inf")) + m = torch.full((h, g * 64), float("-inf")) + l = torch.zeros(h, g * 64) + o = torch.zeros(h, g * 64, vf.size(-1)) + tau = FP8_TAU * math.log(2.0) # FP8_TAU (exp2 units) as natural log + for k0 in range(0, sk, 64): + sb = s[:, :, k0 : k0 + 64] + bm = sb.amax(-1) + exceed = (bm > m + tau) & torch.isfinite(bm) + upd = ( + exceed.view(h, g, 64) + .any(-1, keepdim=True) + .expand(h, g, 64) + .reshape(h, g * 64) + ) + if k0 == 0: + upd = torch.ones_like(upd) + m_new = torch.where(upd, torch.maximum(m, bm), m) + alpha = torch.exp(m - m_new).nan_to_num(0.0) + alpha = torch.where(torch.isfinite(m), alpha, torch.zeros_like(alpha)) + l, o = l * alpha, o * alpha[..., None] + m = m_new + p = torch.exp(sb - m[..., None]).nan_to_num(0.0) + p = ( + torch.where(torch.isfinite(m)[..., None], p, torch.zeros_like(p)) + .to(dtypes.fp8) + .float() + ) + l = l + p.sum(-1) + o = o + torch.einsum("hqk,hkd->hqd", p, vf[:, k0 : k0 + 64]) + o = (o / l[..., None]).nan_to_num(0.0) + lse = torch.where(l > 0, m + torch.log(l), torch.full_like(l, float("-inf"))) + return o[:, :sq].transpose(0, 1), lse[:, :sq] + + +def run_torch_fp8(q8, k8, v8, qd, kd, vd, cu_q, cu_k, *, is_causal, exact=False): + """Reference on the dequantized inputs: out (total_q, hq, 128) fp32, lse + (total_q, hq) fp32, and the exact LSE (total_q, hq). Bottom-right causal; + rows without keys get O = 0, LSE = -inf. Emulates the kernel's online softmax + (see above); exact=True gives the exact softmax instead (O and LSE).""" + total_q, hq, d = q8.shape + hk = k8.size(1) + scale = 1.0 / math.sqrt(d) + out = torch.zeros(total_q, hq, v8.size(-1), dtype=dtypes.fp32) + lse = torch.full((total_q, hq), float("-inf"), dtype=dtypes.fp32) + lse_exact = torch.full((total_q, hq), float("-inf"), dtype=dtypes.fp32) + cuq, cuk = cu_q.tolist(), cu_k.tolist() + for b in range(len(cuq) - 1): + q0, q1, k0, k1 = cuq[b], cuq[b + 1], cuk[b], cuk[b + 1] + sq, sk = q1 - q0, k1 - k0 + if sq == 0 or sk == 0: + continue + qf = q8[q0:q1].float() * qd + kf = (k8[k0:k1].float() * kd).repeat_interleave(hq // hk, dim=1) + vf = (v8[k0:k1].float() * vd).repeat_interleave(hq // hk, dim=1) + s = torch.einsum("qhd,khd->hqk", qf, kf) * scale + if is_causal: + row = torch.arange(sq)[:, None] + col = torch.arange(sk)[None, :] + s = s.masked_fill((col > row + (sk - sq))[None], float("-inf")) + le = torch.logsumexp(s, dim=-1) # (hq, sq) + lse_exact[q0:q1] = le.transpose(0, 1) + if exact: + p = torch.exp(s - le.unsqueeze(-1)).nan_to_num(0.0) + out[q0:q1] = torch.einsum("hqk,khd->qhd", p, vf) + lse[q0:q1] = le.transpose(0, 1) + else: + ob, lb = _emulate_block(s, scale, None, vf.transpose(0, 1)) + out[q0:q1] = ob + lse[q0:q1] = lb.transpose(0, 1) + return out, lse, lse_exact + + +def run_kernel_fp8(q8, k8, v8, qd, kd, vd, cu_q, cu_k, max_q, max_k, *, is_causal, via): + """Returns (out, lse (total_q, nheads)); lse is None for `public`, which + does not return it.""" + scale = 1.0 / math.sqrt(q8.size(-1)) + if via == "public": + out = aiter.flash_attn_varlen_fp8_pertensor_func( + q8, + k8, + v8, + qd, + kd, + vd, + cu_q, + cu_k, + max_q, + max_k, + causal=is_causal, + softmax_scale=scale, + ) + return out, None + if via == "varlen": + r = aiter.flash_attn_varlen_func( + q8, + k8, + v8, + cu_q, + cu_k, + max_q, + max_k, + softmax_scale=scale, + causal=is_causal, + return_lse=True, + q_descale=qd, + k_descale=kd, + v_descale=vd, + ) + return r[0], r[1].transpose(0, 1) + if via == "ops": + out, lse = aiter.fmha_fwd_fp8_varlen_asm( + q8, + k8, + v8, + qd, + kd, + vd, + cu_q, + cu_k, + max_q, + max_k, + softmax_scale=scale, + is_causal=is_causal, + return_lse=True, + ) + return out, lse.squeeze(-1) + raise ValueError(f"unknown via={via!r}") + + +def _fp8_asm_calls(fn): + """Run fn() with a spy on the fp8 ctypes entry; returns the number of launches.""" + mha = aiter.ops.mha + asm = mha._fmha_fwd_fp8_varlen_asm + calls = [] + + def spy(*args, **kwargs): + calls.append(1) + return asm(*args, **kwargs) + + mha._fmha_fwd_fp8_varlen_asm = spy + try: + fn() + finally: + mha._fmha_fwd_fp8_varlen_asm = asm + return len(calls) + + +def _fp8_flops_bytes(seqlens_q, seqlens_k, hq, hk, is_causal): + """2 GEMMs over the visible (q, k) pairs (bottom-right causal); fp8 q/k/v + read once, bf16 O and fp32 LSE written once.""" + + def pairs(sq, sk): + if not is_causal: + return sq * sk + if sk >= sq: + return sq * (sk - sq) + sq * (sq + 1) / 2 + return sk * (sk + 1) / 2 + + flops = sum( + 2.0 * hq * pairs(a, b) * (192 + 128) for a, b in zip(seqlens_q, seqlens_k) + ) + tq, tk = sum(seqlens_q), sum(seqlens_k) + nbytes = tq * hq * 192 + tk * hk * (192 + 128) + tq * hq * (128 * 2 + 4) + return flops, nbytes + + +def check_fp8_against_ref(out, lse, ref, msg, exact=False): + """Asserts the fp8 tolerances above against ref = (O, emulated LSE, exact LSE); + returns the error summary columns. exact: ref's O is the exact softmax (cos + against FP8_COS_MIN_EXACT, as the rounding of P alone moves O from it) and + there is no emulated LSE to check.""" + o, r = out.float(), ref[0] + assert not o.isnan().any(), f"{msg}: NaN in output" + bad = (o - r).abs() > FP8_ATOL + FP8_RTOL * r.abs() + rows = bad.shape[0] * bad.shape[1] + bad_rows = int(bad.any(dim=-1).sum().item()) + if r.abs().max() == 0: # no row sees a key: the output must be exactly 0 + cos = 1.0 if o.abs().max() == 0 else 0.0 + else: + cos = torch.nn.functional.cosine_similarity( + o.flatten(), r.flatten(), dim=0 + ).item() + res = {"max_abs": (o - r).abs().max().item(), "bad rows": bad_rows, "cos": cos} + assert bad_rows <= max(4, rows // 1000), f"{msg}: {bad_rows} of {rows} rows off" + assert cos >= (FP8_COS_MIN_EXACT if exact else FP8_COS_MIN), f"{msg}: cos {cos}" + if lse is not None: + fin = torch.isfinite(ref[2]) + assert torch.equal(torch.isfinite(lse), fin), f"{msg}: LSE finite mask differs" + # a row that sees no key must merge with weight 0: -inf, never NaN / +inf + assert bool( + (lse[~fin] == float("-inf")).all() + ), f"{msg}: LSE of a row without keys is not -inf" + # distance to the exact log-sum-exp: reported, not checked (see above) + res["lse err"] = ( + (lse[fin] - ref[2][fin]).abs().max().item() if fin.any() else 0.0 + ) + if exact: + return res + d = (lse[fin] - ref[1][fin]).abs() + off = int((d > FP8_LSE_EMU_TOL).sum().item()) + assert off <= max( + 4, rows // 1000 + ), f"{msg}: LSE of {off} rows off from the emulated one (max {d.max().item()})" + res["lse vs emu"] = d.max().item() if fin.any() else 0.0 + res["lse rows off"] = off + return res + + +def _record_fp8(ret, name, fn, ref, flops, nbytes, msg): + """Time one candidate and add its us / TFLOPS / TB/s / err columns.""" + (out, lse), us = run_perftest(fn) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = checkAllclose( + ref[0], out.to(dtypes.fp32), rtol=FP8_RTOL, atol=FP8_ATOL, msg=f"{msg} {name}" + ) + return out, lse + + +@benchmark() +def test_fmha_fwd_fp8_varlen_asm(hq, hk, seqlens_q, seqlens_k, is_causal, layout): + """FP8 D192x128: every entry point on dense or strided q/k/v vs the reference; + the strided runs must match the dense one bit for bit.""" + q8, k8, v8, qd, kd, vd, cu_q, cu_k = make_fp8_varlen(seqlens_q, seqlens_k, hq, hk) + max_q, max_k = max(seqlens_q), max(seqlens_k) + ref = run_torch_fp8(q8, k8, v8, qd, kd, vd, cu_q, cu_k, is_causal=is_causal) + dense = run_kernel_fp8( + q8, k8, v8, qd, kd, vd, cu_q, cu_k, max_q, max_k, is_causal=is_causal, via="ops" + ) + if layout != "dense": + q8, k8, v8 = _strided_inputs(layout, q8, k8, v8) + flops, nbytes = _fp8_flops_bytes(seqlens_q, seqlens_k, hq, hk, is_causal) + msg = f"fp8 hq={hq} hk={hk} q={seqlens_q} k={seqlens_k} c={is_causal} {layout}" + + candidates = { + via: functools.partial( + run_kernel_fp8, + q8, + k8, + v8, + qd, + kd, + vd, + cu_q, + cu_k, + max_q, + max_k, + is_causal=is_causal, + via=via, + ) + for via in _FP8_VIAS + } + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + assert _fp8_asm_calls(fn) == 1, f"{msg} via={name}: not one ASM launch" + for _ in range( + _STRIDED_REPEATS + ): # an intermittent race shows up as one differing run + out, lse = fn() + assert torch.equal(out, dense[0]), f"{msg} via={name}: O differs from dense" + assert lse is None or torch.equal( + lse, dense[1] + ), f"{msg} via={name}: LSE differs" + _record_fp8(ret, name, fn, ref, flops, nbytes, msg) + ret.update(check_fp8_against_ref(*dense, ref, msg)) + return ret + + +def merge_attn_states(o_a, lse_a, o_b, lse_b): + """Combine two partial attentions over disjoint key sets (vLLM + merge_attn_states math, fp32). o: (t, h, d); lse: (t, h).""" + m = torch.maximum(lse_a, lse_b) + m = torch.where(torch.isfinite(m), m, torch.zeros_like(m)) + w_a, w_b = torch.exp(lse_a - m), torch.exp(lse_b - m) + s = w_a + w_b + o = (o_a.float() * w_a[..., None] + o_b.float() * w_b[..., None]) / s[..., None] + return o.nan_to_num(0.0), m + torch.log(s) + + +def _gather_tokens(x, cu, lo, hi): + """Pack tokens [start + lo[b], start + hi[b]) of every sequence of x.""" + idx = [torch.arange(s + a, s + b) for s, a, b in zip(cu.tolist()[:-1], lo, hi)] + return x.index_select(0, torch.cat(idx).to(torch.long)) + + +def _kernel_part(q8, k8, v8, qd, kd, vd, cu_q, cu_k, max_q, max_k, causal): + return run_kernel_fp8( + q8, k8, v8, qd, kd, vd, cu_q, cu_k, max_q, max_k, is_causal=causal, via="ops" + ) + + +def _reference_part(q8, k8, v8, qd, kd, vd, cu_q, cu_k, max_q, max_k, causal): + o, lse, _ = run_torch_fp8(q8, k8, v8, qd, kd, vd, cu_q, cu_k, is_causal=causal) + return o, lse + + +def run_chunked_fp8( + q8, + k8, + v8, + qd, + kd, + vd, + cu_q, + cu_k, + prefix_lens, + chunk_lens, + n_ctx, + part=_kernel_part, +): + """MLA-style chunked prefill: the cached prefix in n_ctx non-causal pieces + plus the new tokens causally over themselves, merged with their LSEs. A + piece may hold no key of a sequence; that partial result is O = 0, LSE = -inf. + part=_reference_part runs the same steps on the reference.""" + max_q = max(chunk_lens) + ends = [p + c for p, c in zip(prefix_lens, chunk_lens)] + o, lse = part( + q8, + _gather_tokens(k8, cu_k, prefix_lens, ends), + _gather_tokens(v8, cu_k, prefix_lens, ends), + qd, + kd, + vd, + cu_q, + cu_q, + max_q, + max_q, + True, + ) + for i in range(n_ctx): + lo = [p * i // n_ctx for p in prefix_lens] + hi = [p * (i + 1) // n_ctx for p in prefix_lens] + lens = [b - a for a, b in zip(lo, hi)] + o_p, lse_p = part( + q8, + _gather_tokens(k8, cu_k, lo, hi), + _gather_tokens(v8, cu_k, lo, hi), + qd, + kd, + vd, + cu_q, + _cu(lens), + max_q, + max(max(lens), 1), + False, + ) + o, lse = merge_attn_states(o, lse, o_p, lse_p) + if part is _kernel_part: + o = o.to(dtypes.bf16) # the model keeps the merged output in bf16 + return o, lse + + +@benchmark() +def test_fmha_fwd_fp8_varlen_asm_chunked(hq, hk, prefix_lens, chunk_lens, n_ctx): + """FP8 chunked prefill: one causal call (one_call) vs prefix pieces + new tokens + merged by LSE (merged, the merge in torch so its us is not a kernel time). + exact-prefix (a check, not a candidate) merges the kernel's new-token part with + an exact fp32 prefix part, as when the other part comes from an exact-LSE backend. + References: one_call against the emulated one-call softmax, merged against the + same pieces and merge done on the reference, exact-prefix against the exact + softmax.""" + seqlens_k = [p + c for p, c in zip(prefix_lens, chunk_lens)] + args = make_fp8_varlen(chunk_lens, seqlens_k, hq, hk) + q8, k8, v8, qd, kd, vd, cu_q, cu_k = args + ref_one = run_torch_fp8(*args, is_causal=True) + ref_exact = run_torch_fp8(*args, is_causal=True, exact=True) + ref_merged = ( + *run_chunked_fp8(*args, prefix_lens, chunk_lens, n_ctx, part=_reference_part), + ref_one[2], + ) + refs = {"one_call": ref_one, "merged": ref_merged} + flops, nbytes = _fp8_flops_bytes(chunk_lens, seqlens_k, hq, hk, True) + msg = f"fp8 chunked hq={hq} prefix={prefix_lens} chunk={chunk_lens} n_ctx={n_ctx}" + + candidates = { + "one_call": functools.partial( + run_kernel_fp8, + *args, + max(chunk_lens), + max(seqlens_k), + is_causal=True, + via="ops", + ), + "merged": functools.partial( + run_chunked_fp8, *args, prefix_lens, chunk_lens, n_ctx + ), + } + ret = {"gfx": get_gfx()} + res = {} + for name, fn in candidates.items(): + res[name] = _record_fp8(ret, name, fn, refs[name], flops, nbytes, msg) + for key, val in check_fp8_against_ref( + *res[name], refs[name], f"{msg} {name}" + ).items(): + ret[f"{name} {key}"] = val + ret["merged vs one_call max_abs"] = ( + (res["merged"][0].float() - res["one_call"][0].float()).abs().max().item() + ) + + zeros = [0] * len(prefix_lens) + o_new, lse_new = run_kernel_fp8( + q8, + _gather_tokens(k8, cu_k, prefix_lens, seqlens_k), + _gather_tokens(v8, cu_k, prefix_lens, seqlens_k), + qd, + kd, + vd, + cu_q, + cu_q, + max(chunk_lens), + max(chunk_lens), + is_causal=True, + via="ops", + ) + o_pre, lse_pre, _ = run_torch_fp8( + q8, + _gather_tokens(k8, cu_k, zeros, prefix_lens), + _gather_tokens(v8, cu_k, zeros, prefix_lens), + qd, + kd, + vd, + cu_q, + _cu(prefix_lens), + is_causal=False, + exact=True, + ) + mixed = merge_attn_states(o_new, lse_new, o_pre, lse_pre) + for key, val in check_fp8_against_ref( + *mixed, ref_exact, f"{msg} exact-prefix", exact=True + ).items(): + ret[f"exact-prefix {key}"] = val + return ret + + +# name -> served by the fp8 kernel (_select_gfx1250_fmha_fwd_varlen_asm) +_FP8_GATE_CASES = { + "supported": True, + "per-head descale": False, + "fp64 descale": False, + "no descale": False, + "head dim 128": False, + "bf16 q": False, + "k token stride 2^24 B": False, + "fp32 out": False, + "sink logits": False, + "descale on cpu": False, + "zero kv heads": False, + **{c: ok for c, (_, ok) in _ASM_GATE_OPTION_CASES.items()}, +} + + +@benchmark() +def test_fmha_fwd_fp8_varlen_asm_gate(case): + """The gate every dispatch site asks must pick the fp8 kernel for exactly what + it serves. Not timed: nothing is launched.""" + q8, k8, v8, qd, kd, vd, _, _ = make_fp8_varlen([2], [2], 4, 4) + out, sink = None, None + if case == "per-head descale": + qd = qd.expand(4).contiguous() + elif case == "fp64 descale": + qd = qd.double() + elif case == "no descale": + qd = None + elif case == "head dim 128": + q8, k8 = q8[..., :128], k8[..., :128] + elif case == "bf16 q": + q8 = q8.to(dtypes.bf16) + elif case == "k token stride 2^24 B": + far = torch.empty((1 << 24) + 4 * 192, dtype=k8.dtype) + k8 = far.as_strided(k8.shape, (1 << 24, 192, 1)) + elif case == "fp32 out": + out = torch.empty(2, 4, 128, dtype=dtypes.fp32) + elif case == "sink logits": + sink = torch.zeros(4, dtype=dtypes.fp32) + elif case == "descale on cpu": + qd = qd.cpu() + elif case == "zero kv heads": + k8, v8 = k8[:, :0], v8[:, :0] + mha = aiter.ops.mha + routed = mha._GFX1250_VARLEN_ASM_FP8 == mha._select_gfx1250_fmha_fwd_varlen_asm( + q8, + k8, + v8, + out, + sink_ptr=sink, + q_descale=qd, + k_descale=kd, + v_descale=vd, + **_asm_gate_args(case), + ) + assert routed == _FP8_GATE_CASES[case], f"{case}: gate returned {routed}" + return {"gfx": get_gfx(), "routed to asm": routed} + + +# name -> the C++ entry's error for a direct call with that operand +_FP8_OPS_REFUSED_CASES = { + "descale on cpu": "must be on the same GPU as q", + "zero kv heads": "head counts must be positive", +} + + +@benchmark() +def test_fmha_fwd_fp8_varlen_asm_ops_refused(case): + """A direct call with an operand the kernel cannot use must be refused by the + C++ entry rather than launched. Not timed: nothing is launched.""" + q8, k8, v8, qd, kd, vd, cu_q, cu_k = make_fp8_varlen([2], [2], 4, 4) + if case == "descale on cpu": + qd = qd.cpu() + elif case == "zero kv heads": + k8, v8 = k8[:, :0], v8[:, :0] + try: + aiter.fmha_fwd_fp8_varlen_asm(q8, k8, v8, qd, kd, vd, cu_q, cu_k, 2, 2) + refused = False + except RuntimeError as e: + refused = _FP8_OPS_REFUSED_CASES[case] in str(e) + assert refused, f"{case}: the C++ entry did not refuse it" + return {"gfx": get_gfx(), "ops refused": refused} + + +@benchmark() +def test_fmha_fwd_fp8_varlen_asm_perf(batch, seqlen_q, seqlen_k, hq, hk, is_causal): + """FP8 perf only (the O(s^2) reference is infeasible at these lengths). + + fp8 asm is the kernel under test (public entry point). For comparison: + bf16 asm -- aiter.flash_attn_varlen_func on the same values in bf16 (the bf16 + D192x128 ASM on gfx1250); flydsl fp8 -- FlyDSL's gfx1250 fp8 kernel on the same + fp8 values and descales. TFLOPS use the same FLOP count for all three; TB/s + the fp8 traffic.""" + seqlens_q, seqlens_k = [seqlen_q] * batch, [seqlen_k] * batch + args = make_fp8_varlen(seqlens_q, seqlens_k, hq, hk) + q8, k8, v8, qd, kd, vd, cu_q, cu_k = args + scale = 1.0 / math.sqrt(192) + flops, nbytes = _fp8_flops_bytes(seqlens_q, seqlens_k, hq, hk, is_causal) + q16, k16, v16 = ( + (x.float() * d).to(dtypes.bf16) for x, d in ((q8, qd), (k8, kd), (v8, vd)) + ) + candidates = { + "fp8 asm": functools.partial( + run_kernel_fp8, *args, seqlen_q, seqlen_k, is_causal=is_causal, via="public" + ), + "bf16 asm": functools.partial( + aiter.flash_attn_varlen_func, + q16, + k16, + v16, + cu_q, + cu_k, + seqlen_q, + seqlen_k, + softmax_scale=scale, + causal=is_causal, + ), + } + if flydsl_flash_attn_varlen_fwd is not None: + fly = functools.partial( + flydsl_flash_attn_varlen_fwd, + q8, + k8, + v8, + cu_q, + cu_k, + seqlen_q, + seqlen_k, + softmax_scale=scale, + causal=is_causal, + q_descale=qd, + k_descale=kd, + v_descale=vd, + ) + if fly() is not None: # None: FlyDSL does not take this configuration + candidates["flydsl fp8"] = fly + ret = {"gfx": get_gfx()} + assert ( + _fp8_asm_calls(candidates["fp8 asm"]) == 1 + ), "fp8 asm did not reach the ASM kernel" + for name, fn in candidates.items(): + _, us = run_perftest(fn) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + return ret + + +# (hq, hk, seqlens_q, seqlens_k): one sequence, GQA, q < k (chunked prefill), +# q > k (causal rows that see no key), mixed varlen, short sequences. +_FP8_CORRECTNESS_SHAPES = [ + (16, 16, [1024], [1024]), + (16, 4, [1000], [1000]), + (8, 8, [768], [2048]), + (8, 8, [300], [65]), + (16, 4, [77, 513, 256], [130, 1000, 65]), + (8, 2, [1, 64, 255, 257], [1, 64, 255, 257]), + (32, 32, [2048, 2048], [2048, 2048]), + ( + 8, + 2, + [64, 300], + [0, 0], + ), # no keys at all (a chunked-prefill piece no sequence reaches) + (8, 8, [0, 512, 77], [100, 512, 0]), # a sequence without queries, one without keys +] +# dense plus the strided layouts above (fused_qkv needs equal q and k lengths) +_FP8_LAYOUTS = ["dense"] + _STRIDED_LAYOUTS +# (hq, hk, prefix_lens, chunk_lens, n_ctx); prefix 0 leaves a sequence without +# keys in every prefix piece +_FP8_CHUNKED_SHAPES = [ + (16, 16, [2048], [512], 1), + (16, 16, [3000], [1000], 3), + (8, 2, [1000, 4096, 77], [300, 1024, 129], 2), + (8, 2, [0, 1500, 77], [256, 300, 129], 2), + (8, 2, [0, 0], [256, 129], 1), # the prefix piece holds no key of any sequence + (16, 16, [6144], [2048], 2), +] +# (batch, seqlen_q, seqlen_k, hq, hk); 96 heads with long sequences is the target +# (1 x 12288 / 14336 the typical shape), the rest for reference +_FP8_PERF_SHAPES = [ + (1, 12288, 14336, 96, 96), + (1, 16384, 16384, 96, 96), + (1, 2048, 14336, 96, 96), + (1, 16384, 16384, 32, 32), + (1, 16384, 16384, 64, 4), + (4, 4096, 4096, 128, 128), + (12, 1024, 1024, 128, 128), +] + + def summarize(name, rows): aiter.logger.info( "fmha_fwd_with_sink_varlen_asm %s summary (markdown):\n%s", @@ -628,77 +1434,131 @@ def main(): nargs="*", choices=["randn", "const0.25"], default=["randn", "const0.25"], - help="q/k/v init pattern(s) (default: randn const0.25)", + help="q/k/v init pattern(s) (default: randn const0.25); bf16 only", + ) + parser.add_argument( + "--dtype", + type=str, + nargs="*", + choices=["bf16", "fp8"], + default=["bf16", "fp8"], + help="q/k/v dtype(s): bf16 (fmha_fwd_with_sink_varlen_asm) and/or fp8\n" + "(fmha_fwd_fp8_varlen_asm, D192x128 only) (default: bf16 fp8)", ) args = parser.parse_args() causal_modes = [bool(c) for c in args.causal] - # ---- correctness + perf table ---- - df = [] - for hdim_q, hdim_v, hq, hk, seqlens in _CORRECTNESS_SHAPES: - if hdim_q not in args.head_dim: - continue - for is_causal, init in itertools.product(causal_modes, args.init): - if not is_causal and not _kv_alignment_ok(hdim_q, seqlens): + if "bf16" in args.dtype: + # ---- correctness + perf table ---- + df = [] + for hdim_q, hdim_v, hq, hk, seqlens in _CORRECTNESS_SHAPES: + if hdim_q not in args.head_dim: continue - df.append( - test_fmha_fwd_with_sink_varlen_asm( - hdim_q, hdim_v, hq, hk, seqlens, is_causal, init + for is_causal, init in itertools.product(causal_modes, args.init): + if not is_causal and not _kv_alignment_ok(hdim_q, seqlens): + continue + df.append( + test_fmha_fwd_with_sink_varlen_asm( + hdim_q, hdim_v, hq, hk, seqlens, is_causal, init + ) ) + df = pd.DataFrame(df) + aiter.logger.info( + "fmha_fwd_with_sink_varlen_asm correctness summary (markdown):\n%s", + df.to_markdown(index=False), + ) + summarize( + "bf16 gate", + [test_fmha_fwd_with_sink_varlen_asm_gate(c) for c in _BF16_GATE_CASES], + ) + + # ---- D192x128 strided q/k/v/out ---- + if 192 in args.head_dim: + summarize( + "D192x128 strided q/k/v", + [ + test_fmha_fwd_with_sink_varlen_asm_strided( + hq, hk, seqlens, c, layout + ) + for (hq, hk, seqlens), c, layout in itertools.product( + _STRIDED_SHAPES, causal_modes, _STRIDED_LAYOUTS + ) + ], + ) + summarize( + "D192x128 strided out", + [ + test_fmha_fwd_with_sink_varlen_asm_strided_out( + hq, hk, seqlens, c, w + ) + for (hq, hk, seqlens), c, w in itertools.product( + _STRIDED_SHAPES, causal_modes, _STRIDED_OUT_WIDTHS + ) + ], + ) + summarize( + "D192x128 unsupported stride", + [ + test_fmha_fwd_with_sink_varlen_asm_unsupported_stride(case, c) + for case, c in itertools.product( + ["k token stride 2^24 B", "k/v head stride 0"], causal_modes + ) + ], ) - df = pd.DataFrame(df) - aiter.logger.info( - "fmha_fwd_with_sink_varlen_asm correctness summary (markdown):\n%s", - df.to_markdown(index=False), - ) - # ---- D192x128 strided q/k/v/out ---- - if 192 in args.head_dim: + # ---- perf-only table (large shapes; ref infeasible) ---- + df = [] + for hdim_q, hdim_v, hq, hk, seqlens in _VARLEN_PERF_SHAPES: + if hdim_q not in args.head_dim: + continue + for is_causal, init in itertools.product(causal_modes, args.init): + df.append( + test_fmha_fwd_with_sink_varlen_asm_perf( + hdim_q, hdim_v, hq, hk, seqlens, is_causal, init + ) + ) + df = pd.DataFrame(df) + aiter.logger.info( + "fmha_fwd_with_sink_varlen_asm perf summary (markdown):\n%s", + df.to_markdown(index=False), + ) + + # ---- FP8 D192x128 (fmha_fwd_fp8_varlen_asm) ---- + if "fp8" in args.dtype and 192 in args.head_dim: summarize( - "D192x128 strided q/k/v", + "fp8 D192x128", [ - test_fmha_fwd_with_sink_varlen_asm_strided(hq, hk, seqlens, c, layout) - for (hq, hk, seqlens), c, layout in itertools.product( - _STRIDED_SHAPES, causal_modes, _STRIDED_LAYOUTS + test_fmha_fwd_fp8_varlen_asm(hq, hk, sq, sk, c, layout) + for (hq, hk, sq, sk), c, layout in itertools.product( + _FP8_CORRECTNESS_SHAPES, causal_modes, _FP8_LAYOUTS ) + if layout != "fused_qkv" or sq == sk # one row holds q, k and v ], ) + if True in causal_modes: # chunked prefill is causal + summarize( + "fp8 D192x128 chunked prefill", + [test_fmha_fwd_fp8_varlen_asm_chunked(*s) for s in _FP8_CHUNKED_SHAPES], + ) + summarize( + "fp8 D192x128 gate", + [test_fmha_fwd_fp8_varlen_asm_gate(c) for c in _FP8_GATE_CASES], + ) summarize( - "D192x128 strided out", + "fp8 D192x128 ops refused", [ - test_fmha_fwd_with_sink_varlen_asm_strided_out(hq, hk, seqlens, c, w) - for (hq, hk, seqlens), c, w in itertools.product( - _STRIDED_SHAPES, causal_modes, _STRIDED_OUT_WIDTHS - ) + test_fmha_fwd_fp8_varlen_asm_ops_refused(c) + for c in _FP8_OPS_REFUSED_CASES ], ) summarize( - "D192x128 unsupported stride", + "fp8 D192x128 perf", [ - test_fmha_fwd_with_sink_varlen_asm_unsupported_stride(case, c) - for case, c in itertools.product( - ["k token stride 2^24 B", "k/v head stride 0"], causal_modes - ) + test_fmha_fwd_fp8_varlen_asm_perf(*s, c) + for s, c in itertools.product(_FP8_PERF_SHAPES, causal_modes) ], ) - # ---- perf-only table (large shapes; ref infeasible) ---- - df = [] - for hdim_q, hdim_v, hq, hk, seqlens in _VARLEN_PERF_SHAPES: - if hdim_q not in args.head_dim: - continue - for is_causal, init in itertools.product(causal_modes, args.init): - df.append( - test_fmha_fwd_with_sink_varlen_asm_perf( - hdim_q, hdim_v, hq, hk, seqlens, is_causal, init - ) - ) - df = pd.DataFrame(df) - aiter.logger.info( - "fmha_fwd_with_sink_varlen_asm perf summary (markdown):\n%s", - df.to_markdown(index=False), - ) - if __name__ == "__main__": main() diff --git a/op_tests/test_mha_optional_flydsl.py b/op_tests/test_mha_optional_flydsl.py index c480e72749f..712a3b5466f 100644 --- a/op_tests/test_mha_optional_flydsl.py +++ b/op_tests/test_mha_optional_flydsl.py @@ -71,6 +71,7 @@ def import_module(module_name, globals, locals, fromlist, level): "dtypes": SimpleNamespace(bf16="bf16", fp8="fp8"), "get_gfx": lambda: "gfx1151", "is_gfx1250_asm_supported": lambda: False, + "_IS_GFX1250_ASM_DEVICE": False, } node = _top_level_node(name, ast.FunctionDef) exec( # noqa: S102 - compiles one function parsed from this local module