diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md index 45a42e310528..d671de0bb7c0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/README.md @@ -238,9 +238,9 @@ pass the buffer as `tensors["lse"]`; `dispatch_attention(req).bind_torch(tensors then launches the LSE kernel. A candidate whose kernel cannot write LSE does not declare the `lse` feature, so an LSE request is refused rather than routed to a kernel without the output. The prefill CLI, hipDNN, and the standalone split-KV -decode kernels do not write LSE; unified attention writes it on its split-KV path -only (see `dispatch/attention/README.md`). The gfx942 dense kernel has the same -opt-in. +decode kernels do not write LSE; unified attention writes it on its split-KV and +tiled 2D paths on gfx942/gfx950 (see `dispatch/attention/README.md`). The gfx942 +dense kernel has the same opt-in. Through the dispatcher (opt-in: pin the candidate by `algorithm` and `spec_id`, and optionally a swept point by `tuning_id`): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/README.md b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/README.md index 18a910d81a3a..88e7855be93a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/README.md @@ -305,21 +305,37 @@ Selection pins a spec the same way for dense and unified: another variant is rejected before the search. - `tuning_id="auto"` selects the candidate's base spec. - `emit_lse=True` asks for the optional FP32 LSE output (natural-log softmax - normalizer). It is the `lse` request feature. The dense candidates, the RDNA - WMMA candidates, `attention_unified_3d` and the split-KV tuning candidates - (`*_u3d_*`) declare it; the unified 2D, d256 decode and strided decode - candidates do not. The split-KV support check then refuses LSE on - arches other than gfx942/gfx950 and with sinks. So an LSE request that cannot - be served is refused instead of reaching a kernel without the output. To - launch unified LSE through `bind_torch`, pin a split-KV tuning spec - (`algorithm="unified_tuning"`, - `spec_id="_u3d_..."`): like any unpinned unified request, the - `attention_unified_3d` result routes but cannot be bound. It is a dense output - field (`DENSE_OUTPUT_FIELDS`) and a unified semantic field (`_SEMANTIC_FIELDS`), - so sweeps never vary it, and with it off every spec hash and tuning id is the - one stored before the field existed. The caller passes the buffer to - `bind_torch` as `tensors["lse"]`; the runner validates its shape, dtype, - device, and overlap. The unified layout is `[total_q, Hq, 1]`. + normalizer). It is the `lse` request feature. These declare it: + - the dense candidates; + - the RDNA WMMA candidates; + - `attention_unified_2d` and `attention_unified_3d`; + - every unified tuning candidate (`*_u2d_*`, `*_u3d_*`) except the gfx942 + 4-warp GQA one (`gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm`). + + The d256 decode, strided decode and route-only `attention_gfx942_dense_pipe` + / `attention_gfx950_d256` candidates do not. Each path's own support check + (split-KV or tiled 2D) refuses LSE on arches other than gfx942/gfx950 and with + sinks. The scalar fallback writes no LSE, so it refuses it. When the tiled 2D + check declines an LSE request that `select_path()` sends to tiled 2D (gfx942 + D256 prefill at block size 32 exceeds its LDS budget), + `attention_unified_3d` serves it with split-KV, as the runner's `auto` backend + does (`select_auto_path`). An LSE request that cannot be served is refused + instead of reaching a kernel without the output. + + To launch unified LSE through `bind_torch`, pin a tuning spec + (`algorithm="unified_tuning"`, `spec_id="_u2d_..."` or `"_u3d_..."`). + Like any unpinned unified request, the `attention_unified_2d` / `_3d` result + routes but cannot be bound. + + `emit_lse` is a dense output field (`DENSE_OUTPUT_FIELDS`) and a unified + semantic field (`_SEMANTIC_FIELDS`), so sweeps never vary it. It is kept out + of the config defaults fingerprint, so a pin's `tuning_id` is the same with + LSE on or off. With it off, every spec hash and tuning id is the one stored + before the field existed. + + The caller passes the buffer to `bind_torch` as `tensors["lse"]`; the runner + validates its shape, dtype, device and overlap. The unified layout is + `[total_q, Hq, 1]`. A candidate resolves a request once: `support` and `select_spec` share a per-candidate cache, so a pin is not resolved twice. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/bindings.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/bindings.py index 6b1d391e5ee8..3a7ead406796 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/bindings.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/bindings.py @@ -108,6 +108,13 @@ def validate_tuning_attention_contract(request, problem, spec) -> None: raise ValueError("kernel spec K/V storage dtype disagrees with problem") if bool(spec.fp8_fnuz) != bool(problem.fp8_fnuz): raise ValueError("tuning wrapper FP8 encoding disagrees with problem") + if spec.path == "2d": + spec_lse = bool(getattr(kernel_spec, "emit_lse", False)) + if spec_lse != problem.emit_lse: + raise ValueError( + f"kernel_spec.emit_lse={spec_lse} disagrees with " + f"problem.emit_lse={problem.emit_lse}" + ) reduce_spec = spec.reduce_spec if reduce_spec is not None: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py index 394993df9f29..3e256ae4fdaf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py @@ -10,6 +10,8 @@ from __future__ import annotations +import dataclasses + from typing import Callable, Mapping, Optional from kernels.common.attention_unified import ( @@ -44,11 +46,10 @@ def _tuning_precheck(variant: AttentionGeometryVariant): - """Coverage for one geometry variant. The scalar coverage check knows nothing - about the LSE output, so an LSE request also has to pass the gate of the - variant's own path: the split-KV gate owns which arches and features can - write LSE, and the 2D gate refuses it. That holds even if a 2D variant ever - declared the ``lse`` feature by mistake.""" + """Coverage for one geometry variant. The scalar coverage check refuses LSE + (the scalar kernel writes none), so it is asked about the LSE-off problem + and an LSE request must also pass the gate of the variant's own path: the + split-KV or tiled 2D gate owns which arches and features can write LSE.""" path_gate = ( supports_native_unified_attention_3d_tiled if variant.path == "3d" @@ -57,8 +58,12 @@ def _tuning_precheck(variant: AttentionGeometryVariant): def precheck(req): problem = _tuning_problem(req) - ok, why = supports_native_unified_attention(problem, arch=variant.arch) - if ok and problem.emit_lse: + if not problem.emit_lse: + return supports_native_unified_attention(problem, arch=variant.arch) + ok, why = supports_native_unified_attention( + dataclasses.replace(problem, emit_lse=False), arch=variant.arch + ) + if ok: return path_gate(problem, arch=variant.arch) return ok, why @@ -88,6 +93,7 @@ def signature(spec): include_bt_stride=True, include_qq_bias_stride=True, kv_dtype=spec.kernel_spec.kv_storage_dtype, + emit_lse=bool(getattr(spec.kernel_spec, "emit_lse", False)), ) def bind_torch(request, spec, tensors, **kwargs): @@ -112,11 +118,12 @@ def bind_torch(request, spec, tensors, **kwargs): ShapeRange("hdim_q", allowed=UNIFIED_HEAD_SIZES), ShapeRange("kv_block_size", allowed=UNIFIED_BLOCK_SIZES), ), - # Only split-KV variants write LSE; 2D variants refuse the feature. + # Every unified variant writes LSE except the gfx942 4-warp GQA + # builder (see _GFX942_4WARP_LSE_REFUSAL). supports_features=( - ATTENTION_FEATURES | {"lse"} - if variant.path == "3d" - else ATTENTION_FEATURES + ATTENTION_FEATURES + if variant.builder_kind == "gfx942_4warp_gqa" + else ATTENTION_FEATURES | {"lse"} ), ), space=unified_space(variant), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py index f2ddac05f769..5daf630deba7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py @@ -201,7 +201,7 @@ class AttentionRequest(OperatorRequest): # Also write the natural-log softmax normalizer (LSE) into a caller-owned # FP32 ``tensors["lse"]``. Reported as the ``lse`` feature, so only # candidates whose kernels implement it (dense, RDNA WMMA, and the - # gfx942/gfx950 unified split-KV path) admit it. + # gfx942/gfx950 unified split-KV and tiled 2D paths) admit it. emit_lse: bool = False # Serve runtime head widths from one compiled bucket: ``hdim_q``/``hdim_v`` # are the bucket maxima, and any widths the kernel takes up to them (RDNA diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/generic.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/generic.py index 46722879c8bc..bc4f3eb106fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/generic.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/generic.py @@ -15,8 +15,10 @@ from typing import Tuple from kernels.common.attention_unified import ( + select_auto_path, supports_native_unified_attention, supports_native_unified_attention_3d_tiled, + supports_native_unified_attention_tiled, ) from rocke.core.arch import known_arches from rocke.dispatch.core import ( @@ -63,9 +65,9 @@ # the scalar kernel elsewhere (tiled gates reject it off gfx950). supports_features=ATTENTION_FEATURES, ) -# Only the split-KV (3D) path writes the optional LSE output; its support check -# (the 3D gate below) still refuses it on arches/features that cannot. -_UNIFIED_3D_CAPABILITY = dataclasses.replace( +# Both unified paths write the optional LSE output; each path's own gate (in +# ``support`` below) still refuses it on arches/features that cannot. +_UNIFIED_LSE_CAPABILITY = dataclasses.replace( _UNIFIED_CAPABILITY, supports_features=ATTENTION_FEATURES | {"lse"} ) @@ -83,17 +85,34 @@ def support(req: OperatorRequest) -> Tuple[bool, str]: if not ok: return False, why problem = _problem(req) - ok, why = supports_native_unified_attention(problem, arch=req.arch) + # The scalar gate is the coverage check; it refuses LSE (the scalar + # kernel writes none), so ask it about the LSE-off problem and let the + # path's own gate rule on LSE below. + coverage = ( + dataclasses.replace(problem, emit_lse=False) + if problem.emit_lse + else problem + ) + ok, why = supports_native_unified_attention(coverage, arch=req.arch) if not ok: return False, why - if problem.select_path(req.arch) != path: + # Split-KV also takes the LSE requests that ``select_path(arch)`` sends + # to tiled 2D but whose tiled 2D gate declines, as auto does + # (``select_auto_path``); the tiled 2D candidate keeps its gate's reason. + if problem.select_path(req.arch) != path and not ( + path == "3d" and select_auto_path(problem, req.arch) == "3d" + ): return False, ( f"problem routes to {problem.select_path(req.arch)!r} path, not {path!r}" ) if problem.emit_lse: - # The scalar coverage check above knows nothing about LSE; the - # split-KV gate owns which arches and features can write it. - ok, why = supports_native_unified_attention_3d_tiled(problem, arch=req.arch) + # The path's own gate owns which arches and features can write LSE. + gate = ( + supports_native_unified_attention_3d_tiled + if path == "3d" + else supports_native_unified_attention_tiled + ) + ok, why = gate(problem, arch=req.arch) if not ok: return False, why return True, "ok" @@ -126,7 +145,7 @@ def select(req: OperatorRequest) -> AttentionSpec: spec_id=spec_id, abi_version=ATTENTION_ABI_VERSION, priority=priority, - capability=_UNIFIED_3D_CAPABILITY if path == "3d" else _UNIFIED_CAPABILITY, + capability=_UNIFIED_LSE_CAPABILITY, _supports=support, select_spec=select, signature=lambda _spec: (), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py index cd093fc05840..d2f733268899 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/tuning_specs.py @@ -13,6 +13,7 @@ from typing import Mapping, Optional, Tuple from kernels.common.attention_unified import ( + _GFX942_4WARP_LSE_REFUSAL, UnifiedAttentionProblem, _kv_block_bytes, _kv_cache_policy_3d, @@ -208,6 +209,8 @@ def make_explicit_attention_2d_spec( validate_explicit_fp8_encoding( arch=arch, use_fp8=problem.use_fp8, fp8_fnuz=problem.fp8_fnuz ) + if problem.emit_lse and config.builder_kind == "gfx942_4warp_gqa": + raise ValueError(_GFX942_4WARP_LSE_REFUSAL) if config.num_warps not in NUM_WARPS: raise ValueError(f"num_warps must be one of {NUM_WARPS}") if config.block_m_per_warp not in BLOCK_M_PER_WARP: @@ -231,6 +234,9 @@ def make_explicit_attention_2d_spec( block_m_per_warp=int(config.block_m_per_warp), tile_size=int(tile_size), waves_per_eu=config.waves_per_eu, + # The output the caller asked for, never a tuning knob (see + # _SEMANTIC_FIELDS); passed only when on, so LSE-off specs are as before. + **({"emit_lse": True} if problem.emit_lse else {}), **knobs, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/unified_rules.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/unified_rules.py index 383ee2455f10..a81cc2d56a9e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/unified_rules.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/unified_rules.py @@ -389,14 +389,15 @@ def defaults(self, base, kernel): recorded = self.recorded(base) return { "codepath": dict(self.fixed(base)), + # ``emit_lse`` is an output the problem selects, not a default that + # knobs are a delta against: keeping it out of both the 2D kernel + # and the split-KV reduce defaults leaves every config_key, and so + # every stored pin, as before the field existed. "kernel": { name: value for name, value in _declared_defaults(type(kernel.kernel_spec)) - if name not in recorded + if name not in recorded and name != "emit_lse" }, - # ``emit_lse`` is an output the problem selects, not a default that - # knobs are a delta against: keeping it out leaves every split-KV - # config_key, and so every stored pin, as before the field existed. "reduce": ( None if reduce is None diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_lse_store.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_lse_store.py index 9662f24d83aa..dee472727d2e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_lse_store.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_lse_store.py @@ -9,9 +9,10 @@ LSE store, ``dense_row_has_key`` / ``zero_keyless_row`` the empty-row output masking, and ``validate_dense_lse`` is the host-side contract for the caller-owned LSE buffer; -* the unified split-KV reduce: the same IR helpers emit its LSE store, and - ``validate_unified_lse`` is its host-side buffer contract. Both validators - sit on the common ``_check_lse_buffer``; +* the unified split-KV reduce and the unified tiled 2D epilogues (gfx942 / + gfx950): the same IR helpers emit their LSE store, ``lse_row_has_mass`` the + ``!(l == 0)`` row test, and ``validate_unified_lse`` is their host-side buffer + contract. Both validators sit on the common ``_check_lse_buffer``; * the MFMA / WMMA attention forward inner bodies: ``make_*_lse_epilogue`` produce *epilogue callables* for :class:`rocke.helpers.attention_fwd_ext.AttnFwdExt.epilogue_hook`. @@ -61,6 +62,7 @@ "LSE_ROWS_PER_LANE", "dense_row_has_key", "lse_row_value", + "lse_row_has_mass", "lse_store", "make_lse_epilogue", "make_mfma_lse_epilogue", @@ -90,6 +92,17 @@ def lse_store( b.global_store(dst, element_offset, value, align=4) +def lse_row_has_mass(b: IRBuilder, l: Value) -> Value: + """``!(l == 0)``, NaN included: the exact inverse of the tiled 2D O guard + ``fcmp("oeq", l, 0)``. A keyless row (l == 0) gets -inf; a NaN row keeps NaN. + + Built as ``one OR uno`` instead of ``lnot(oeq)``: the HIP lowerer maps i1 to + ``bool`` and lowers ``arith.not`` to ``~x``, which is true for both inputs. + """ + zero = b.const_f32(0.0) + return b.lor(b.fcmp("one", l, zero), b.fcmp("uno", l, zero)) + + def dense_row_has_key( b, query_row, kv_length, *, causal, diagonal, window, right_bound=0 ): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 7119ff3575e2..a519efb1f07e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -148,8 +148,9 @@ class UnifiedAttentionProblem: causal_top_left: bool = False right_bound: int = 0 # Optional FP32 natural-log softmax normalizer, [total_q, num_query_heads, 1]. - # Implemented by the gfx942/gfx950 split-KV (3D) reduce only; every other - # path declines it (see _reject_lse). + # Implemented by the gfx942/gfx950 split-KV (3D) reduce and tiled 2D + # kernels (not the gfx942 4-warp GQA builder); every other path declines it + # (see _reject_lse). emit_lse: bool = False def __post_init__(self): @@ -560,29 +561,39 @@ def _reject_mask_on_arch( return None -_LSE_SPLIT_KV_ARCHES = ("gfx942", "gfx950") +_LSE_ARCHES = ("gfx942", "gfx950") +_LSE_PATHS = {"3d": "split-KV", "tiled 2D": "tiled 2D"} def _reject_lse( problem: UnifiedAttentionProblem, arch: str, path: str ) -> Optional[Tuple[bool, str]]: - """The LSE output exists only in the gfx942/gfx950 split-KV reduce. + """The LSE output exists in the gfx942/gfx950 split-KV reduce and tiled 2D + kernels only. - ``path`` is ``"3d"`` for split-KV, else a label for the refused path.""" + ``path`` is ``"3d"`` (split-KV) or ``"tiled 2D"``; any other label names a + path that writes no LSE and is refused.""" if not problem.emit_lse: return None - if path != "3d": + label = _LSE_PATHS.get(path) + if label is None: return False, f"LSE output is not implemented on the {path} attention path" - if arch not in _LSE_SPLIT_KV_ARCHES: - return False, f"LSE output is not implemented on {arch} split-KV attention" + if arch not in _LSE_ARCHES: + return False, f"LSE output is not implemented on {arch} {label} attention" if problem.use_sinks: return False, ( - "LSE output with attention sinks is not implemented on split-KV " + "LSE output with attention sinks is not implemented on unified " "attention (keyless rows would drop the sink)" ) return None +_GFX942_4WARP_LSE_REFUSAL = ( + "LSE output is not implemented on the gfx942 4-warp GQA builder (finite " + "-1e30 max sentinel, no l == 0 guard, fixed 17-argument ABI)" +) + + def supports_native_unified_attention( problem: UnifiedAttentionProblem, arch: str, @@ -607,6 +618,11 @@ def supports_native_unified_attention( - softcap: yes """ arch = _canonical_arch_for_selection(arch) + # The scalar kernel writes no LSE; refusing here keeps the auto tail from + # ever launching it for an LSE request. + rejected = _reject_lse(problem, arch, "scalar 2D") + if rejected is not None: + return rejected if problem.head_size not in UNIFIED_HEAD_SIZES: return False, f"unsupported head_size {problem.head_size}" if problem.block_size not in UNIFIED_BLOCK_SIZES: @@ -1100,7 +1116,7 @@ class _TiledRoute: def _gfx942_4warp_eligible(problem: "UnifiedAttentionProblem", arch: str) -> bool: - """Shared eligibility gate for the gfx942 4-warp GQA cohorts -- the eight + """Shared eligibility gate for the gfx942 4-warp GQA cohorts -- the clauses common to ``_d256_gfx942_fast`` and ``_d128_gfx942_swa_fast``. Extracting them here dedupes the two AND-walls; each cohort predicate adds only its distinctive clauses (head_size / dtype / sliding_window / @@ -1116,6 +1132,10 @@ def _gfx942_4warp_eligible(problem: "UnifiedAttentionProblem", arch: str) -> boo # shipped default builder). Exclude caches > 2 GiB, which need the i64 # path -- they fall back to the correct default builder. and not _enable_i64_kv_addr(problem) + # The 4-warp GQA bodies write no LSE (finite -1e30 max sentinel, no + # l == 0 guard, fixed 17-argument ABI): an LSE request takes the + # default gfx942 builder, which does. + and not problem.emit_lse ) @@ -1731,7 +1751,7 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: num_warps_2d = None if fp16_flash else _select_2d_num_warps(problem, arch) tile_size = _select_2d_tile_size(problem, arch) block_m_per_warp = _select_2d_block_m_per_warp(problem, arch) - return ( + key = ( "tiled", arch, _binary_search_iters(problem), @@ -1870,6 +1890,11 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: else None ), ) + _mask_key(problem) + if problem.emit_lse: + # Suffix only when on, so every LSE-off key is unchanged (same as the + # 3D key). The 4-warp exclusion is already keyed by _gfx942_4warp_fast. + key = key + ("lse",) + return key def _select_2d_waves_per_eu( @@ -3390,6 +3415,9 @@ def _tiled_spec_from_problem( _spec = _impl(problem, arch) if overrides is not None: _spec = replace(_spec, **overrides) + if problem.emit_lse: + # Output contract, not a tuning choice; LSE-off specs stay as before. + _spec = replace(_spec, emit_lse=True) return _resolve_lds_budget(_spec, arch) @@ -4304,6 +4332,7 @@ def _attn_signature( include_qq_bias_stride: bool = False, include_additive_bias: bool = False, kv_dtype: Optional[str] = None, + emit_lse: bool = False, ): from rocke.helpers.spec import SignatureBuilder @@ -4335,9 +4364,13 @@ def _attn_signature( sb.scalar("block_table_stride", "i32") if include_qq_bias_stride: sb.scalar("qq_bias_stride_0", "i32") - # Appended last, matching the kernel's conditional param declarations. + # Appended last, matching the kernel's conditional param declarations: + # additive bias, then the LSE output (the tiled 2D builders declare + # lse_ptr after declare_additive_bias_params). if include_additive_bias: append_additive_bias_signature(sb) + if emit_lse: + sb.ptr("lse_ptr", "f32") return sb.build() @@ -4368,6 +4401,7 @@ def _attn_values( k_scale: float = 1.0, v_scale: float = 1.0, out_scale: float = 1.0, + lse=None, kv_addr_limit: _KvAddrLimit | None, ): _require_additive_bias(problem, additive_bias) @@ -4404,6 +4438,8 @@ def _attn_values( vals["additive_bias_batch_stride"] = int(additive_bias_batch_stride) vals["additive_bias_head_stride"] = int(additive_bias_head_stride) vals["additive_bias_sq_stride"] = int(additive_bias_sq_stride) + if problem.emit_lse: + vals["lse_ptr"] = lse return vals @@ -5074,7 +5110,8 @@ def _explicit_path_supported( """ # LSE eligibility is an output contract, not a coverage gate: the # allow_unsupported escape hatch must not skip it (a keyless row would - # silently drop an attention sink, or a 2D path would write no LSE at all). + # silently drop an attention sink, or an arch without the store would + # write no LSE at all). rejected = _reject_lse( problem, _canonical_arch_for_selection(arch), @@ -5097,6 +5134,26 @@ def _explicit_path_supported( raise ValueError(f"unknown attention path kind {kind!r}") +def select_auto_path(problem: UnifiedAttentionProblem, arch: str) -> str: + """The path ``auto`` tries for ``problem`` on ``arch``: ``select_path(arch)``, + except an LSE request the tiled 2D gate declines goes to split-KV. + + Only split-KV and tiled 2D write LSE (the scalar fallback writes none), so + when tiled 2D cannot take a 2D-routed LSE request, split-KV is the one path + left that can. Example: gfx942 D256 bf16 prefill at block size 32, whose + LSE-off form runs on the 4-warp builder (no LSE store) and whose default + tiled spec exceeds the LDS budget. LSE off, this is ``select_path(arch)``. + """ + path = problem.select_path(arch) + if ( + path == "2d" + and problem.emit_lse + and not _explicit_path_supported(problem, None, "2d", arch)[0] + ): + return "3d" + return path + + def _recommend_graph_replay(problem: UnifiedAttentionProblem) -> bool: """Graph-vs-ungraph heuristic. @@ -5686,6 +5743,8 @@ def _get_2d_launcher( if tuning_spec is not None else _kv_storage_dtype(problem) ), + # The runner has already checked that an explicit spec agrees. + emit_lse=problem.emit_lse, ), cache_key=("2d",) + cache_key, ) @@ -6044,9 +6103,11 @@ def run_unified_attention_torch( ``lse`` is the optional caller-owned natural-log softmax normalizer, a contiguous FP32 ``[total_q, num_query_heads, 1]`` tensor, written only when - ``problem.emit_lse`` is set. Only the gfx942/gfx950 split-KV (3D) path - writes it; every other path raises ``NotImplementedError`` rather than - returning without it. Rows with no visible key get ``-inf`` (and O = 0). + ``problem.emit_lse`` is set. The gfx942/gfx950 split-KV (3D) and tiled 2D + paths write it (an LSE request skips the gfx942 2D graph replay and the + 4-warp GQA builder); the scalar path, sinks and other arches raise + ``NotImplementedError`` rather than returning without it. Rows with no + visible key get ``-inf`` (and O = 0). """ # Caller-owned LSE output ([total_q, Hq, 1] FP32, ``problem.emit_lse``): # validated before any compile so a bad buffer never reaches a launch. @@ -6246,9 +6307,9 @@ def run_unified_attention_torch( # Explicit specs are initially selected before framework tensors exist. # Refresh address-width state and tuning identity from the real cache. tuning_spec = tuning_spec.with_num_kv_blocks(int(k.shape[0])) - # A path that cannot write LSE (a 2D spec, or a split-KV reduce without - # the output, e.g. gfx1250) is refused as NotImplementedError, the same - # error as the backend refusals below. + # A path that cannot write LSE (gfx1250, sinks, the gfx942 4-warp GQA + # builder) is refused as NotImplementedError, the same error as the + # backend refusals below. if problem.emit_lse: rejected = _reject_lse( problem, @@ -6257,14 +6318,19 @@ def run_unified_attention_torch( ) if rejected is not None: raise NotImplementedError(rejected[1]) - # Otherwise the explicit spec owns the reduce kernel, so it must agree - # with the problem on the LSE output: a mismatch would either drop the - # output or launch a kernel whose argument list differs from the - # signature. A reduce spec without the field cannot emit LSE. - reduce_lse = tuning_spec.path == "3d" and bool( - getattr(tuning_spec.reduce_spec, "emit_lse", False) + if getattr(tuning_spec, "builder_kind", None) == "gfx942_4warp_gqa": + raise NotImplementedError(_GFX942_4WARP_LSE_REFUSAL) + # Otherwise the explicit spec owns the kernel that writes LSE (the 3D + # reduce, or the 2D kernel itself), so it must agree with the problem: + # a mismatch would either drop the output or launch a kernel whose + # argument list differs from the signature. A spec without the field + # cannot emit LSE. + lse_owner = ( + tuning_spec.reduce_spec + if tuning_spec.path == "3d" + else tuning_spec.kernel_spec ) - if reduce_lse != problem.emit_lse: + if bool(getattr(lse_owner, "emit_lse", False)) != problem.emit_lse: raise ValueError("tuning_spec LSE output disagrees with problem.emit_lse") # Auto path selection. Historically we *always* preferred 3D when @@ -6289,8 +6355,16 @@ def run_unified_attention_torch( # buckets (where our old auto picked 3D for chunked prefill at 30- # 150× the right path's cost) while keeping the decode-shape 3D # wins on the parity harness (those shapes still satisfy the "3D - # is fine" branch of ``use_2d_kernel``). - prefer_2d = backend == "auto" and problem.select_path(_launch_arch) == "2d" + # is fine" branch of ``use_2d_kernel``). An LSE request the tiled 2D gate + # declines goes to split-KV, the only other path that writes LSE + # (``select_auto_path``); an explicit spec keeps ``select_path(arch)``. + auto_path = ( + problem.select_path(_launch_arch) + if tuning_spec is not None + else select_auto_path(problem, _launch_arch) + ) + prefer_2d = backend == "auto" and auto_path == "2d" + reason_3d = None # split-KV verdict, named by the auto LSE refusal below # FP8 K/V prefill that no tiled 2D kernel serves (gfx942 has no FP8 2D # path) runs on the 3D split-KV kernel, which handles multi-token query # blocks, instead of the scalar fallback. use_fp8 is checked first so other @@ -6340,37 +6414,6 @@ def run_unified_attention_torch( if backend == "3d": raise NotImplementedError(reason_3d) - if problem.emit_lse: - # Every remaining path (2D graph replay, tiled 2D, scalar) writes no LSE. - # Refuse here, BEFORE the 2D graph fast path: _run_2d_graphed replays a - # cache hit before any support gate, and _cheap_2d_sig does not encode - # emit_lse, so an LSE-off graph captured on the same buffers would - # otherwise serve this request and leave lse unwritten. In order: a - # split-KV LSE reason (sinks, arch); for ``auto``, the split-KV gate's - # verdict (name the backend='3d' way out only when that gate admits the - # problem); otherwise the explicitly requested 2D path. - rejected = _reject_lse(problem, _launch_arch, "3d") - if rejected is not None: - raise NotImplementedError(rejected[1]) - if backend == "auto": - if prefer_2d: - ok_3d, reason_3d = _explicit_path_supported( - problem, tuning_spec, "3d", _launch_arch - ) - if ok_3d: - raise NotImplementedError( - "LSE output is not implemented on the tiled 2D attention " - "path, which backend='auto' chose for this shape; the " - "split-KV path can serve it: pass backend='3d' or pin a " - "'*_u3d_*' tuning spec (it can be slower for this shape)" - ) - raise NotImplementedError( - "LSE output needs the split-KV attention path, which declined " - f"this problem: {reason_3d}" - ) - _, why = _reject_lse(problem, _launch_arch, "tiled 2D") - raise NotImplementedError(why) - if backend in ("tiled", "auto"): # Graphed fast path FIRST: for the short-context regime the launch is # auto-graphed (heuristic: _recommend_graph_replay). The graph is looked @@ -6378,8 +6421,12 @@ def run_unified_attention_torch( # a replay skips supports + cache_key + the kernarg pack -- the host # overhead that otherwise dominates tiny-shape latency. Skipped when the # caller is already capturing the forward (they take precedence). + # Never for LSE: _run_2d_graphed replays a cache hit before any gate, + # and _cheap_2d_sig / the graph key / its refs carry no LSE, so an + # LSE-off graph on the same buffers would serve the request unwritten. if ( tuning_spec is None + and not problem.emit_lse and _enable_2d_graph_replay(problem, _launch_arch) and not _torch_stream_capturing() ): @@ -6459,6 +6506,7 @@ def run_unified_attention_torch( v_scale=v_scale, out_scale=out_scale, kv_addr_limit=meta.kv_addr_limit, + lse=lse, ) return launcher( vals, @@ -6471,6 +6519,17 @@ def run_unified_attention_torch( if backend == "tiled": raise NotImplementedError(reason_t) + if problem.emit_lse and backend == "auto": + # Only split-KV and tiled 2D write LSE; the scalar fallback below does + # not (its gate refuses LSE), so name why each LSE path declined. + declined = f"tiled 2D: {reason_t}" + if reason_3d is not None: + declined = f"split-KV: {reason_3d}; {declined}" + raise NotImplementedError( + "LSE output needs the split-KV or tiled 2D attention path, which " + f"declined this problem ({declined})" + ) + # Scalar fallback. Uses the same KernelLauncher infrastructure as # the tiled paths so module load + arg lifetime + stream resolution # are handled by-construction. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index bd11aa58ac0a..6f094c5c56eb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -34,10 +34,12 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field import os from typing import Optional, Tuple +from kernels.common._lse_store import lse_row_has_mass, lse_store + from rocke.core.ir import ( BF16, CACHE_POLICY_NAMES, @@ -478,8 +480,16 @@ class UnifiedAttention2DTiledSpec: # only when the kernel is VALU/throughput-bound (the x8 conflict-free-V # path); no-op-guarded to sliding_window==0. use_causal_mask_phase_split: bool = False + # emit_lse: also write the natural-log softmax normalizer (LSE) as FP32 + # [total_q, num_query_heads, 1] through a trailing ``lse_ptr`` argument. + # Off by default; off emits no new IR, so shipped kernels are unchanged. + emit_lse: bool = field(default=False, kw_only=True) def __post_init__(self): + if type(self.emit_lse) is not bool: + raise ValueError( + f"emit_lse must be a bool, got {type(self.emit_lse).__name__}" + ) # gfx942 (CDNA3) variant: the narrow ``16x16x16`` default path only. # The wide-K (32x32x16 migration) and transpose-read / native-fp8 # experimental knobs all depend on gfx950-only ISA (wide-K MFMA atoms, @@ -1050,6 +1060,7 @@ def kernel_name(self) -> str: "iglp1" if self.use_iglp_opt else "", "cmps" if self.use_causal_mask_phase_split else "", "k1buf" if self.use_k_single_buffer else "", + "lse" if self.emit_lse else "", ) @@ -1569,6 +1580,14 @@ def build_unified_attention_2d_tiled( num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) qq_bias_stride0_p = b.param("qq_bias_stride_0", I32) + # Appended LAST and only when requested, so the LSE-off ABI is unchanged. + lse = ( + b.param( + "lse_ptr", PtrType(F32, "global"), noalias=True, writeonly=True, align=4 + ) + if spec.emit_lse + else None + ) if spec.use_q_major_grid: q_block_global_idx = b.block_id_x() @@ -5504,6 +5523,12 @@ def _strided_v_b_operand(k_iter: int) -> Value: final = kvloop.results l_final = [final[2 * r + 1] for r in range(SOFTMAX_STATE_SLOTS)] + # Loop carry is (m0, l0, m1, l1, ..., acc...); m is in log2 units + # (qk_scale = scale * log2(e)). Read only for LSE so the LSE-off IR is + # unchanged. + m_final = ( + [final[2 * r] for r in range(SOFTMAX_STATE_SLOTS)] if spec.emit_lse else None + ) ml_count_final = 2 * SOFTMAX_STATE_SLOTS # acc_final indexed by (n * ACC_M_ATOMS + atom) acc_final = [ @@ -5520,6 +5545,30 @@ def _acc_final_get(n: int, atom: int) -> Value: rcp_l = [b.rcp(l_final[r]) for r in range(SOFTMAX_STATE_SLOTS)] l_zero = [b.fcmp("oeq", l_final[r], zero_f) for r in range(SOFTMAX_STATE_SLOTS)] + def _store_row_lse(row: Value, slot: int, writer: Value) -> None: + """Write LSE for CTA row ``row`` from softmax state ``slot``. + + Same token/head/validity mapping as the O stores below. ``writer`` + picks one lane among the lanes that hold a replicated copy of the + row's (m, l). ``!(l == 0)`` is the row-has-mass test, the inverse of + the O guard: all-masked rows end with l == 0, a zero-trip loop with + m == -inf (lse_store writes -inf for both), and a NaN l gives NaN LSE. + """ + pos = b.add(qb_start_pos, b.div(row, b.const_i32(NQK))) + qh = b.add(b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK))) + ok = b.land(b.cmp_lt(pos, cur_batch_q_len), b.cmp_lt(qh, b.const_i32(NUM_QH))) + if VALID_ROWS < BLOCK_M: + ok = b.land(ok, b.cmp_lt(row, b.const_i32(VALID_ROWS))) + with b.scf_if(b.land(writer, ok)): + lse_store( + b, + lse, + b.add(b.mul(b.add(cu_q_start, pos), b.const_i32(NUM_QH)), qh), + m_final[slot], + l_final[slot], + row_has_mass=lse_row_has_mass(b, l_final[slot]), + ) + if USE_MFMA_32X32: if TRANSPOSED_QK_32X32: # Per-lane direct scalar global stores. Tested both this and a @@ -5554,6 +5603,25 @@ def _acc_final_get(n: int, atom: int) -> Value: ) inv_l_t = b.rcp(l_final[0]) l_zero_t = b.fcmp("oeq", l_final[0], zero_f) + if spec.emit_lse: + # One query row per lane (lane % 32); lanes l and l^32 hold the + # same state (xor-32 reductions). The lower half writes. Slot 0: + # scalar state has 1 slot; without it all 16 slots are + # identical and O also reads slot 0. + # Same operand order as gfx950: the C fastkv_regp wrapper builds + # gfx950 kernels through this arch's C epilogue. + with b.scf_if(b.land(op_mask_t, b.cmp_eq(lane_half32, b.const_i32(0)))): + lse_store( + b, + lse, + b.add( + b.mul(b.add(cu_q_start, op_pos_t), b.const_i32(NUM_QH)), + op_qh_t, + ), + m_final[0], + l_final[0], + row_has_mass=lse_row_has_mass(b, l_final[0]), + ) for n in range(ACC_N_TILES): acc32 = _acc_final_get(n, 0) for reg in range(REGS_PER_LANE): @@ -5605,6 +5673,17 @@ def _acc_final_get(n: int, atom: int) -> Value: head=op_qh32_base, dim=b.const_i32(0), ) + if spec.emit_lse: + # Slot r is C-row wave_row_base + _mfma_32x32_c_row(lane, r); the + # 32 lanes of a half share it (_warp_xor_reduce_sum_32lane). Lanes + # 0 and 32 write; together they cover the warp's 32 rows. These + # stores come before the first Acc_lds b.sync(), so the barrier + # pattern is unchanged. + writer32 = b.cmp_eq(lane_col32, b.const_i32(0)) + for r in range(SOFTMAX_STATE_SLOTS): + _store_row_lse( + b.add(wave_row_base, _mfma_32x32_c_row(b, lane, r)), r, writer32 + ) for n in range(ACC_N_TILES): acc32 = _acc_final_get(n, 0) for reg in range(REGS_PER_LANE): @@ -5742,6 +5821,17 @@ def _acc_final_get(n: int, atom: int) -> Value: if stripe + 1 < OUT_STRIPES: b.sync() + if spec.emit_lse: + # After the O stripes, not before: emitted ahead of them, the LSE store + # let LLVM fuse the O multiply and FP16 convert into v_fma_mixlo_f16 + # (one rounding; LSE-off rounds twice), moving FP16 O by 1 ulp versus + # LSE-off on MI300X/MI355X (ROCm 10). Here the O ISA matches LSE-off. + # Slot r is row wave_row_base + _in_warp_row(r), shared by the 16 + # lane_col lanes (_warp_xor_reduce_sum). Lanes 0/16/32/48 + # (lane_col == 0) write; each owns REGS_PER_LANE rows. + writer16 = b.cmp_eq(lane_col, b.const_i32(0)) + for r in range(SOFTMAX_STATE_SLOTS): + _store_row_lse(b.add(wave_row_base, _in_warp_row(r)), r, writer16) return b.kernel @@ -6072,8 +6162,13 @@ def build_gfx942_4warp_gqa( ) -> KernelDef: """Emit the gfx942 4-warp GQA D256 paged-attention ``KernelDef``.""" from ..common.attention_arch import require_tiled_attention_arch + from ..common.attention_unified import _GFX942_4WARP_LSE_REFUSAL require_tiled_attention_arch(arch) + if spec.emit_lse: + # This body has no LSE store; refuse rather than build a kernel that + # carries the "lse" name but takes no lse_ptr and writes nothing. + raise NotImplementedError(_GFX942_4WARP_LSE_REFUSAL) if spec.use_i64_kv_addr: # K/V are i32 element offsets. Refuse rather than emit an i32 kernel # under a spec that claims i64; dispatch keeps the flag off through diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py index 5756ed33970e..0b3debf3624d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py @@ -38,7 +38,7 @@ strided_kv_offset, strided_kv_resource, ) -from kernels.common._lse_store import lse_store +from kernels.common._lse_store import lse_row_has_mass, lse_store from rocke.core.ir import ( BF16, @@ -1488,7 +1488,9 @@ def build_unified_attention_reduce_tiled( # Every lane holds the same butterfly-reduced overall_max/expsum, so one # lane writes. overall_max is in log2 units (segm_max carries # scale*log2e); lse_store converts to natural log and writes -inf for a - # row with no mass (every segment empty), whose O is already 0. + # row with no mass (every segment empty, sum 0), whose O is already 0. + # A NaN sum (NaN query row) keeps NaN, as O does: the same !(l == 0) + # test the tiled 2D epilogues use, never l > 0. with b.scf_if(b.cmp_eq(tid, b.const_i32(0))): lse_index = b.add(b.mul(q_token, b.const_i32(NUM_QH)), q_head) lse_store( @@ -1497,7 +1499,7 @@ def build_unified_attention_reduce_tiled( lse_index, overall_max, overall_expsum, - row_has_mass=b.fcmp("ogt", overall_expsum, zero_f), + row_has_mass=lse_row_has_mass(b, overall_expsum), ) return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index 0d794dcfb818..2a6218f45d97 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -39,13 +39,14 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional, Tuple from kernels.common._attention_additive_bias import ( add_additive_bias, declare_additive_bias_params, ) +from kernels.common._lse_store import lse_row_has_mass, lse_store from rocke.core.ir import ( BF16, @@ -500,8 +501,16 @@ class UnifiedAttention2DTiledSpec: # to the alignment chosen by ``causal_top_left``. Not combinable with ALiBi / # QQ-bias, whose offsets are defined on the causal diagonal. right_bound: int = 0 + # emit_lse: also write the natural-log softmax normalizer (LSE) as FP32 + # [total_q, num_query_heads, 1] through a trailing ``lse_ptr`` argument. + # Off by default; off emits no new IR, so shipped kernels are unchanged. + emit_lse: bool = field(default=False, kw_only=True) def __post_init__(self): + if type(self.emit_lse) is not bool: + raise ValueError( + f"emit_lse must be a bool, got {type(self.emit_lse).__name__}" + ) if type(self.use_additive_bias) is not bool: raise ValueError( "use_additive_bias must be a bool, got " @@ -962,6 +971,7 @@ def kernel_name(self) -> str: "fp8pv" if self.use_fp8_mfma_pv else "", "regpv" if self.use_register_pv else "", "qrev" if self.use_qblock_reverse else "", + "lse" if self.emit_lse else "", ) @@ -1306,6 +1316,17 @@ def build_unified_attention_2d_tiled( # Additive-bias params are appended only when the feature is on, so the # non-bias kernel ABI (and its IR) is unchanged. additive_bias = declare_additive_bias_params(b) if USE_ADDITIVE_BIAS else None + lse = ( + b.param( + "lse_ptr", + PtrType(F32, "global"), + noalias=True, + writeonly=True, + align=4, + ) + if spec.emit_lse + else None + ) kv_head_idx = b.block_id_x() q_block_global_idx = b.block_id_y() @@ -4535,6 +4556,9 @@ def _apply_transposed_pv_regs(acc32: Value, n: int, p_regs) -> Value: final = kvloop.results l_final = [final[2 * r + 1] for r in range(SOFTMAX_STATE_SLOTS)] + if spec.emit_lse: + # m is in log2 units (scores carry scale*log2e), reduced across lanes. + m_final = [final[2 * r] for r in range(SOFTMAX_STATE_SLOTS)] ml_count_final = 2 * SOFTMAX_STATE_SLOTS # acc_final indexed by (n * ACC_M_ATOMS + atom) acc_final = [ @@ -4585,6 +4609,21 @@ def _acc_final_get(n: int, atom: int) -> Value: ) inv_l_t = b.rcp(l_final[0]) l_zero_t = b.fcmp("oeq", l_final[0], zero_f) + if spec.emit_lse: + # One row per lane; lanes l and l^32 hold the same (m, l), so + # only the low half writes. Emitted once, outside the n-loop. + with b.scf_if(b.land(op_mask_t, b.cmp_eq(lane_half32, b.const_i32(0)))): + lse_store( + b, + lse, + b.add( + b.mul(b.add(cu_q_start, op_pos_t), b.const_i32(NUM_QH)), + op_qh_t, + ), + m_final[0], + l_final[0], + row_has_mass=lse_row_has_mass(b, l_final[0]), + ) for n in range(ACC_N_TILES): acc32 = _acc_final_get(n, 0) for reg in range(REGS_PER_LANE): @@ -4636,6 +4675,27 @@ def _acc_final_get(n: int, atom: int) -> Value: head=op_qh32_base, dim=b.const_i32(0), ) + if spec.emit_lse: + # State slot reg's row is wave_row_base + _mfma_32x32_c_row(lane, reg) + # (= hoist_row[reg]); all 32 lanes of a half agree, lane_col32 == 0 + # (lanes 0 and 32) writes. + lse_writer32 = b.cmp_eq(lane_col32, b.const_i32(0)) + for reg in range(REGS_PER_LANE): + with b.scf_if(b.land(lse_writer32, hoist_row_ok[reg])): + lse_store( + b, + lse, + b.add( + b.mul( + b.add(cu_q_start, hoist_qp_r[reg]), + b.const_i32(NUM_QH), + ), + hoist_qh_r[reg], + ), + m_final[reg], + l_final[reg], + row_has_mass=lse_row_has_mass(b, l_final[reg]), + ) for n in range(ACC_N_TILES): acc32 = _acc_final_get(n, 0) for reg in range(REGS_PER_LANE): @@ -4735,7 +4795,6 @@ def _acc_final_get(n: int, atom: int) -> Value: head=op_qh, dim=b.const_i32(0), ) - for stripe in range(OUT_STRIPES): n_start = stripe * N_TILES_PER_STRIPE # ---- stage 1: write this stripe's N-tiles into Acc_lds ---- @@ -4773,6 +4832,31 @@ def _acc_final_get(n: int, atom: int) -> Value: if stripe + 1 < OUT_STRIPES: b.sync() + if spec.emit_lse: + # After the O stripes, not before: emitted ahead of them, the LSE store + # let LLVM fuse the O multiply and FP16 convert into v_fma_mixlo_f16 + # (one rounding; LSE-off rounds twice), moving FP16 O by 1 ulp versus + # LSE-off on MI300X/MI355X (ROCm 10). Here the O ISA matches LSE-off. + # Slot r's row is wave_row_base + _in_warp_row(r) (= hoist_row[r]); the + # 16 lanes of a row group agree, lane_col == 0 (one per lane_rg) writes. + lse_writer16 = b.cmp_eq(lane_col, b.const_i32(0)) + for reg in range(REGS_PER_LANE): + with b.scf_if(b.land(lse_writer16, hoist_row_ok[reg])): + lse_store( + b, + lse, + b.add( + b.mul( + b.add(cu_q_start, hoist_qp_r[reg]), + b.const_i32(NUM_QH), + ), + hoist_qh_r[reg], + ), + m_final[reg], + l_final[reg], + row_has_mass=lse_row_has_mass(b, l_final[reg]), + ) + return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py index 3fd81037e687..12217b92b2a8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py @@ -39,7 +39,7 @@ strided_kv_offset, strided_kv_resource, ) -from kernels.common._lse_store import lse_store +from kernels.common._lse_store import lse_row_has_mass, lse_store from rocke.core.ir import ( BF16, @@ -1651,7 +1651,9 @@ def build_unified_attention_reduce_tiled( # Every lane holds the same butterfly-reduced overall_max/expsum, so one # lane writes. overall_max is in log2 units (segm_max carries # scale*log2e); lse_store converts to natural log and writes -inf for a - # row with no mass (every segment empty), whose O is already 0. + # row with no mass (every segment empty, sum 0), whose O is already 0. + # A NaN sum (NaN query row) keeps NaN, as O does: the same !(l == 0) + # test the tiled 2D epilogues use, never l > 0. with b.scf_if(b.cmp_eq(tid, b.const_i32(0))): lse_index = b.add(b.mul(q_token, b.const_i32(NUM_QH)), q_head) lse_store( @@ -1660,7 +1662,7 @@ def build_unified_attention_reduce_tiled( lse_index, overall_max, overall_expsum, - row_has_mass=b.fcmp("ogt", overall_expsum, zero_f), + row_has_mass=lse_row_has_mass(b, overall_expsum), ) return b.kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py index 8fb0b1ebd547..f0b9486acf4f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_attention.py @@ -22,10 +22,11 @@ # is listed here, so a widening like fp8 can never silently reach a path that # does not implement it. # ``lse`` is an output feature, not part of ATTENTION_FEATURES: only candidates -# whose kernels write LSE (dense, RDNA WMMA, ``attention_unified_3d`` and the -# ``_u3d_`` split-KV tuning candidates) declare it, so every other candidate -# must keep declining it. The generated ``_u3d_`` candidates are not in this map; -# their feature set is checked in the tuning-candidate test below. +# whose kernels write LSE (dense, RDNA WMMA, ``attention_unified_2d``, +# ``attention_unified_3d`` and every unified tuning candidate except the gfx942 +# 4-warp GQA one) declare it, so every other candidate must keep declining it. +# The generated tuning candidates are not in this map; their feature set is +# checked in the tuning-candidate test below. _GFX950_DENSE_FEATURES = { "band", "causal", @@ -102,6 +103,7 @@ "fp8", "sinks", "sliding_window", + "lse", }, "attention_unified_3d": { "band", @@ -111,7 +113,7 @@ "fp8", "sinks", "sliding_window", - # Split-KV only: the 2D path writes no LSE (see attention_unified_2d). + # Both unified paths write LSE (tiled 2D and split-KV). "lse", }, } @@ -397,9 +399,9 @@ def test_declared_features_are_frozen(self): tuning = [c for c in attention_candidates() if c.algorithm == "unified_tuning"] self.assertTrue(tuning) for candidate in tuning: - # Only split-KV (u3d) geometries write the optional LSE output. + # Every unified geometry writes LSE except the gfx942 4-warp GQA one. expected = set(ATTENTION_FEATURES) | ( - {"lse"} if "_u3d_" in candidate.name else set() + set() if "_u2d_gfx942_4warp_" in candidate.name else {"lse"} ) with self.subTest(candidate=candidate.name): self.assertEqual(set(candidate.capability.supports_features), expected) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_lse_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_lse_dispatch.py index b9b8ed38cc11..6e2e75936efe 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_lse_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_dense_lse_dispatch.py @@ -56,13 +56,15 @@ def test_emit_lse_rejects_non_boolean_values(bad_emit_lse): @pytest.mark.parametrize("arch", sorted(_DENSE_SPEC_IDS)) -def test_unpinned_lse_request_is_refused_not_served_without_lse(arch): - """This 512/512 prefill routes to the unified 2D path, which writes no LSE - (only unified split-KV does), so ``auto`` must fail loudly instead of - returning a kernel that silently drops the requested output.""" - dispatch_attention(_request(arch)) # the same problem without LSE routes - with pytest.raises(ValueError, match="lse"): - dispatch_attention(_request(arch, emit_lse=True)) +def test_unpinned_lse_request_is_served_with_lse(arch): + """This 512/512 prefill routes to the unified 2D path, which writes LSE on + gfx942/gfx950, so ``auto`` serves the LSE request with the LSE variant of + the same kernel, never one that silently drops the requested output.""" + off = dispatch_attention(_request(arch)) + on = dispatch_attention(_request(arch, emit_lse=True)) + assert on.kernel_id.candidate == off.kernel_id.candidate + assert on.spec.emit_lse is True + assert on.spec.kernel_name() == off.spec.kernel_name() + "_lse" @pytest.mark.parametrize("arch,spec_id", _CASES) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx1151_attention_selection.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx1151_attention_selection.py index d0e0139a523a..64392777671f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx1151_attention_selection.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx1151_attention_selection.py @@ -414,9 +414,13 @@ def test_lse_reuses_the_same_code_object(self): plain = dispatch_attention(make(layout=layout)).spec lse = dispatch_attention(make(layout=layout, emit_lse=True)).spec self.assertEqual(plain, lse) + # On CDNA the same request routes to unified tiled 2D, whose LSE + # variant is a different binary (one more pointer argument). for arch in ("gfx950", "gfx942"): - with self.subTest(arch=arch), self.assertRaises(ValueError): - dispatch_attention(_request(arch=arch, emit_lse=True)) + with self.subTest(arch=arch): + plain = dispatch_attention(_request(arch=arch)).spec + lse = dispatch_attention(_request(arch=arch, emit_lse=True)).spec + self.assertEqual(lse.kernel_name(), plain.kernel_name() + "_lse") def test_score_features_layouts_and_fp8_reach_both_kernels(self): features = dict( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py index f6f4a2f055dc..05ad82a158b2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py @@ -1,7 +1,7 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Dispatcher contract for the optional unified split-KV LSE output (host only).""" +"""Dispatcher contract for the optional unified (split-KV and tiled 2D) LSE output (host only).""" import dataclasses @@ -48,6 +48,15 @@ def _pinned(req, spec_id): return dataclasses.replace(req, algorithm=_TUNING_ALGORITHM, spec_id=spec_id) +def _prefill_request(arch, **overrides): + """Single-sequence 256-token prefill: routes to tiled 2D.""" + return _decode_request(arch, batch=1, seqlen_q=256, seqlen_k=256, **overrides) + + +_NARROW_2D_SPEC = "{arch}_u2d_narrow_nw1_mw16_t1xb_llvm" +_GFX942_4WARP_SPEC = "gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm" + + @pytest.mark.parametrize("arch", _ARCHES) def test_unpinned_lse_decode_request_routes_to_split_kv(arch): result = dispatch_attention(_decode_request(arch, emit_lse=True)) @@ -87,7 +96,7 @@ def test_pinned_split_kv_spec_follows_the_request(arch): # What a long-lived cache stored for ``_decode_request`` before emit_lse existed: # the unpinned spec_hash, and the split-KV pin with its spec_hash. These are -# develop's own values (gfx950's split-KV config_key moved with #13120's +# develop's own values (gfx950's split-KV config_key moved with the # additive-bias field); emit_lse must leave them exactly as develop computes them. # Both split-KV config_keys moved again once decode kept K/V cacheable when it # fits the last-level cache: this shape's 64 MiB of K/V now records @@ -114,34 +123,150 @@ def test_pre_lse_cache_entries_still_resolve_with_lse_off(arch): assert dispatch_attention(pinned).kernel_id.spec_hash == spec_hash +# What a long-lived cache stored for ``_prefill_request`` before the 2D specs +# had emit_lse: unpinned spec_hash, narrow pin + hash. Develop's own values +# (gfx950's pin moved with the additive-bias field, then with use_qblock_reverse). +_PRE_LSE_2D_UNPINNED_SPEC_HASH = "c7dd16a4a52afb9d" +_PRE_LSE_2D_PIN = { + "gfx942": ( + "narrow_nw1_mw16_t1xb_llvm_wpenone@98141a2e646271d9", + "738d09721910c8ad", + ), + "gfx950": ( + "narrow_nw1_mw16_t1xb_llvm_wpenone@7a7cfcc8a5d56e43", + "c97e33ec12b3b1a4", + ), +} + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_pre_lse_2d_cache_entries_still_resolve_with_lse_off(arch): + unpinned = dispatch_attention(_prefill_request(arch)) + assert unpinned.kernel_id.spec_hash == _PRE_LSE_2D_UNPINNED_SPEC_HASH + tuning_id, spec_hash = _PRE_LSE_2D_PIN[arch] + pinned = dataclasses.replace( + _pinned(_prefill_request(arch), _NARROW_2D_SPEC.format(arch=arch)), + tuning_id=tuning_id, + ) + assert dispatch_attention(pinned).kernel_id.spec_hash == spec_hash + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_unpinned_lse_prefill_request_routes_to_tiled_2d(arch): + on = dispatch_attention(_prefill_request(arch, emit_lse=True)) + off = dispatch_attention(_prefill_request(arch)) + assert on.spec.path == "2d" and on.spec.emit_lse is True + assert on.spec.kernel_name() == off.spec.kernel_name() + "_lse" + assert on.kernel_id.spec_hash != off.kernel_id.spec_hash + + @pytest.mark.parametrize("arch", _ARCHES) @pytest.mark.parametrize("emit_lse", [False, True]) def test_split_kv_sweep_keeps_the_requested_output(arch, emit_lse): """A tuning sweep varies performance knobs only; it must never add or drop - the LSE output. With LSE requested, no 2D configuration may be offered.""" + the LSE output, on either path.""" results = dispatch_attention_all( _decode_request(arch, emit_lse=emit_lse), candidate_prefix="attention_" ) split_kv = [r for r in results if r.spec.path == "3d"] assert split_kv, "no split-KV configuration was enumerated" assert {r.spec.reduce_spec.emit_lse for r in split_kv} == {emit_lse} + two_d = [r for r in results if r.spec.path == "2d"] + assert {r.spec.kernel_spec.emit_lse for r in two_d} <= {emit_lse} if emit_lse: - assert not [r for r in results if r.spec.path == "2d"] + assert not [r for r in two_d if r.spec.builder_kind == "gfx942_4warp_gqa"] + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_pinned_2d_spec_follows_the_request(arch): + spec_id = _NARROW_2D_SPEC.format(arch=arch) + on = dispatch_attention(_pinned(_prefill_request(arch, emit_lse=True), spec_id)) + off = dispatch_attention(_pinned(_prefill_request(arch), spec_id)) + assert on.spec.kernel_spec.emit_lse is True + assert off.spec.kernel_spec.emit_lse is False + # Same configuration (emit_lse is an output, not a knob), distinct binary. + assert on.spec.config_key == off.spec.config_key + assert on.spec.kernel_name() == off.spec.kernel_name() + "_lse" + assert on.spec.cache_key() != off.spec.cache_key() @pytest.mark.parametrize("arch", _ARCHES) -def test_pinned_2d_spec_refuses_lse(arch): - prefill = _decode_request(arch, batch=1, seqlen_q=256, seqlen_k=256) - spec_id = f"{arch}_u2d_narrow_nw1_mw16_t1xb_llvm" - dispatch_attention(_pinned(prefill, spec_id)) # served without LSE - # Refused by the missing ``lse`` feature or, should a 2D variant ever declare - # it, by the 2D gate ("LSE output is not implemented ..."): either is a refusal. +def test_2d_lse_sweep_carries_the_output(arch): + results = dispatch_attention_all( + _prefill_request(arch, emit_lse=True), candidate_prefix="attention_" + ) + two_d = [r for r in results if r.spec.path == "2d"] + assert two_d, "no tiled 2D configuration was enumerated" + assert {r.spec.kernel_spec.emit_lse for r in two_d} == {True} + + +def test_pinned_gfx942_4warp_spec_refuses_lse(): + req = _prefill_request("gfx942", hdim_q=256, hdim_v=256) with pytest.raises(ValueError, match="(?i)lse"): dispatch_attention( - _pinned(dataclasses.replace(prefill, emit_lse=True), spec_id) + _pinned(dataclasses.replace(req, emit_lse=True), _GFX942_4WARP_SPEC) ) +def test_explicit_4warp_spec_builder_refuses_lse(): + from dispatch.attention.common import _tuning_problem + from dispatch.attention.tuning_specs import ( + ExplicitAttention2DConfig, + make_explicit_attention_2d_spec, + ) + + problem = _tuning_problem( + _prefill_request("gfx942", hdim_q=256, hdim_v=256, emit_lse=True) + ) + config = ExplicitAttention2DConfig( + num_warps=1, + block_m_per_warp=32, + tile_policy="64", + builder_kind="gfx942_4warp_gqa", + ) + with pytest.raises(ValueError, match="4-warp GQA builder"): + make_explicit_attention_2d_spec(problem, config, arch="gfx942") + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_lse_with_sinks_is_refused_on_2d(arch): + with pytest.raises(ValueError, match="LSE output with attention sinks"): + dispatch_attention(_prefill_request(arch, emit_lse=True, use_sinks=True)) + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_binding_rejects_a_2d_kernel_spec_that_disagrees_on_lse(arch): + from dispatch.attention.bindings import validate_tuning_attention_contract + from dispatch.attention.common import _tuning_problem + + request = _prefill_request(arch, emit_lse=True) + result = dispatch_attention(_pinned(request, _NARROW_2D_SPEC.format(arch=arch))) + spec = dataclasses.replace( + result.spec, + kernel_spec=dataclasses.replace(result.spec.kernel_spec, emit_lse=False), + ) + with pytest.raises(ValueError, match="kernel_spec.emit_lse"): + validate_tuning_attention_contract(request, _tuning_problem(request), spec) + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_binding_forwards_the_lse_buffer_on_2d(arch, monkeypatch): + from dispatch.attention.common import _tuning_problem + from kernels.common import attention_unified as au + + seen = {} + monkeypatch.setattr(au, "run_unified_attention_torch", lambda **kw: seen.update(kw)) + request = _pinned( + _prefill_request(arch, emit_lse=True), _NARROW_2D_SPEC.format(arch=arch) + ) + problem = _tuning_problem(request) + lse = object() + tensors = {**_paged_tensors(problem), "lse": lse} + tensors["cu_seqlens_q"] = _FakeTensor((2,), "torch.int32", [0, problem.total_q]) + dispatch_attention(request).bind_torch(tensors).launch() + assert seen["lse"] is lse and seen["backend"] == "tiled" + + @pytest.mark.parametrize("arch", _ARCHES) @pytest.mark.parametrize("pin", [False, True], ids=["auto", "pinned"]) def test_lse_with_sinks_is_refused(arch, pin): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_domain.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_domain.py new file mode 100644 index 000000000000..3dc039002a5d --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_domain.py @@ -0,0 +1,326 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""CPU differential test of unified LSE dispatch over arch x path x feature. + +Each key runs LSE off and on in the same process. A key served both ways keeps +its candidate, path, tuning id and spec; only ``emit_lse`` and the ``lse`` +kernel-name token change, apart from the named reroutes in ``_REROUTES``. LSE +on: a served request runs a variant that writes LSE, and every unified path +that refuses it says why in its own words. Host only. +""" + +import ast +import dataclasses +import itertools +import re + +import pytest + +from dispatch.attention import ( + AttentionRequest, + attention_candidates, + dispatch_attention, +) +from dispatch.attention.common import _problem +from kernels.common.attention_unified import select_auto_path +from rocke.helpers.spec import kernel_name_join + +_ARCHES = ("gfx942", "gfx950", "gfx90a", "gfx1250") +_LSE_ARCHES = ("gfx942", "gfx950") +_CUS = {"gfx942": 304, "gfx950": 256} +_SHAPES = { # name: (seqlen_q, seqlen_k, mask_type, sliding_window, kv_block_size) + "decode_long": (1, 4096, 2, 0, 16), + "decode_short": (1, 256, 2, 0, 16), + "decode_swa": (1, 4096, 2, 128, 16), + "prefill": (2048, 2048, 1, 0, 16), + # Block size 32: gfx942 D256 bf16 LSE no longer fits the tiled 2D LDS + # budget, so auto hands it to split-KV. + "prefill_bs32": (2048, 2048, 1, 0, 32), +} +# Pinned unified tuning specs: the shape that pins each, and its id per arch. +_PINS = { + "split_kv": ( + "decode_long", + {a: f"{a}_u3d_splitkv_seg16_t1xb" for a in _LSE_ARCHES}, + ), + "u2d_narrow": ( + "prefill", + {a: f"{a}_u2d_narrow_nw1_mw16_t1xb_llvm" for a in _LSE_ARCHES}, + ), + "u2d_4warp": ("prefill", {"gfx942": "gfx942_u2d_gfx942_4warp_nw1_mw32_t64_llvm"}), +} +_SINKS = ( + "LSE output with attention sinks is not implemented on unified attention " + "(keyless rows would drop the sink)" +) +_LABEL = {"3d": "split-KV", "2d": "tiled 2D"} +_UNIFIED = re.compile(r"attention_(unified_(2d|3d)|gfx9\d\d_u(2d|3d)_\w+)$") +_NO_FEATURES = "capability: cannot serve features " +# LSE requests that tiled 2D declines (LDS budget) and auto hands to split-KV. +# LSE off, bf16 runs the 4-warp builder and fp16 the scalar kernel; neither +# writes LSE. +_HANDOVER = frozenset( + f"gfx942|prefill_bs32|d256|{dt}|sinks=0|paged|auto" for dt in ("bf16", "fp16") +) +# Candidates that do not declare the ``lse`` feature, by name. +_NO_LSE = frozenset( + c.name + for c in attention_candidates() + if "lse" not in c.capability.supports_features +) + + +def _domain(): + for arch, shape, hd, dt, sinks, lse, kv, pin in itertools.product( + _ARCHES, + _SHAPES, + (128, 256), + ("bf16", "fp16"), + (False, True), + (False, True), + ("paged", "strided"), + ("auto", *_PINS), + ): + if pin != "auto" and ( + _PINS[pin][0] != shape or arch not in _PINS[pin][1] or kv != "paged" + ): + continue + sq, sk, mask_type, window, block = _SHAPES[shape] + fields = dict( + batch=4, nhead_q=32, nhead_k=8, seqlen_q=sq, seqlen_k=sk, + hdim_q=hd, hdim_v=hd, arch=arch, mask_type=mask_type, + sliding_window=window, dtype=dt, use_sinks=sinks, emit_lse=lse, + kv_layout=kv, num_cus=_CUS.get(arch, 128), kv_block_size=block, + ) # fmt: skip + if pin != "auto": + fields.update(algorithm="unified_tuning", spec_id=_PINS[pin][1][arch]) + key = f"{arch}|{shape}|d{hd}|{dt}|sinks={int(sinks)}|{kv}|{pin}" + yield key, lse, AttentionRequest(**fields) + + +def _run(request): + """(outcome row, result or None): the served candidate and path, or the + refusal message.""" + try: + result = dispatch_attention(request) + except ValueError as exc: + return {"ok": False, "why": str(exc)}, None + return { + "ok": True, + "candidate": result.kernel_id.candidate, + "path": result.spec.path, + }, result + + +def _unified_verdicts(request): + """{candidate: (admitted, reason)} for every unified candidate.""" + return { + c.name: c.admits(request) + for c in attention_candidates() + if _UNIFIED.match(c.name) + } + + +def _writes_lse(result): + spec = result.spec + if not hasattr(spec, "kernel_spec"): # unpinned unified AttentionSpec + return spec.emit_lse and spec.kernel_name().endswith("_lse") + if spec.path == "3d": + return spec.reduce_spec.emit_lse + return spec.kernel_spec.emit_lse and spec.builder_kind != "gfx942_4warp_gqa" + + +def _expected_lse_refusal(candidate, arch, sinks, off_why): + """The reason a unified candidate must give when LSE makes it refuse, or + None when LSE must not change its verdict. ``off_why`` is its LSE-off + reason: the capability check lists every missing feature at once.""" + match = _UNIFIED.match(candidate) + path = match.group(2) or match.group(3) + if "_u2d_gfx942_4warp_" in candidate: + missing = {"lse"} + if off_why.startswith(_NO_FEATURES): + missing |= set(ast.literal_eval(off_why[len(_NO_FEATURES) :])) + return _NO_FEATURES + str(sorted(missing)) + if arch not in _LSE_ARCHES: + return f"LSE output is not implemented on {arch} {_LABEL[path]} attention" + if sinks: + return _SINKS + return None + + +@pytest.fixture(scope="module") +def table(): + on, off = {}, {} + for key, lse, request in _domain(): + (on if lse else off)[key] = ( + request, + *_run(request), + _unified_verdicts(request), + ) + return on, off + + +# The only keys whose LSE-on dispatch may leave the LSE-off candidate. +_REROUTES = ("tiled_2d_lse_handover", "no_lse_candidate") + + +def _reroute(off, request): + """(name, candidate or None for any LSE writer, path) of the reroute LSE on + must take from the LSE-off result ``off``, or None when it must not move.""" + if ( + off.spec.path == "2d" + and select_auto_path(_problem(request), request.arch) == "3d" + ): + # The tiled 2D gate declines the LSE problem, so auto hands it to split-KV. + return "tiled_2d_lse_handover", "attention_unified_3d", "3d" + if off.kernel_id.candidate in _NO_LSE: + # Writes no LSE (route-only dense, d256/strided decode, gfx942 4-warp), so + # dispatch falls through to the next candidate on the same path that does. + return "no_lse_candidate", None, off.spec.path + return None + + +def _field_diffs(off, on, prefix=""): + """Dotted names of spec fields that differ, recursing into nested specs and + skipping ``emit_lse``.""" + if dataclasses.is_dataclass(off) and type(off) is type(on): + return [ + name + for f in dataclasses.fields(off) + if f.name != "emit_lse" + for name in _field_diffs( + getattr(off, f.name), getattr(on, f.name), f"{prefix}{f.name}." + ) + ] + return [] if off == on else [prefix.rstrip(".") or type(off).__name__] + + +def _lse_flags(spec): + """Every ``emit_lse`` value in ``spec`` and its nested specs.""" + if not dataclasses.is_dataclass(spec): + return [] + return [ + flag + for f in dataclasses.fields(spec) + for flag in ( + [getattr(spec, f.name)] + if f.name == "emit_lse" + else _lse_flags(getattr(spec, f.name)) + ) + ] + + +def _kernel_names(spec): + """{kernel: name} for the kernels ``spec`` launches, and which one writes + LSE: a pinned split-KV reduce, else the main kernel.""" + names = {"kernel": spec.kernel_name()} + reduce = getattr(spec, "reduce_spec", None) + if spec.path == "3d" and reduce is not None: + names["reduce"] = reduce.kernel_name() + return names, "reduce" if "reduce" in names else "kernel" + + +def _lse_on_identity_diff(off, on, request): + """Why the LSE-on dispatch is not the LSE-off one plus LSE, or None.""" + reroute = _reroute(off, request) + if reroute: + name, want_candidate, want_path = reroute + got = (on.kernel_id.candidate, on.spec.path) + if got[1] != want_path or want_candidate not in (None, got[0]): + return f"{name}: wants {(want_candidate, want_path)}, got {got}" + return None if _writes_lse(on) else f"{name}: {got} writes no LSE" + off_id = (off.kernel_id.candidate, off.spec.path, off.kernel_id.tuning_id) + on_id = (on.kernel_id.candidate, on.spec.path, on.kernel_id.tuning_id) + if on_id != off_id: + return f"unlisted reroute {off_id} -> {on_id}" + names, writer = _kernel_names(off.spec) + want = {**names, writer: kernel_name_join(names[writer], "lse")} + if _kernel_names(on.spec)[0] != want: + return f"kernel names {_kernel_names(on.spec)[0]}, want {want}" + fields = _field_diffs(off.spec, on.spec) + if fields: + return f"spec fields differ: {fields}" + if any(_lse_flags(off.spec)) or not any(_lse_flags(on.spec)): + return f"emit_lse off={_lse_flags(off.spec)} on={_lse_flags(on.spec)}" + return None + + +def test_lse_on_identity_is_lse_off_identity_plus_lse(table): + """Per key, in one run: LSE on serves the LSE-off identity plus LSE, or a + named ``_REROUTES`` entry. A key LSE off refuses stays refused; keys LSE on + refuses are the per-path wording test's.""" + on, off = table + assert on.keys() == off.keys() + bad, fired = {}, set() + for key, (request, row, result, _) in on.items(): + off_row, off_result = off[key][1], off[key][2] + if not off_row["ok"]: + if row["ok"]: + bad[key] = f"LSE on serves {row}, LSE off refuses" + continue + if not row["ok"]: + continue + reroute = _reroute(off_result, request) + if reroute: + fired.add(reroute[0]) + why = _lse_on_identity_diff(off_result, result, request) + if why: + bad[key] = why + assert not bad, bad + # Each reroute must still describe a key in the domain, or it is dead. + assert fired == set(_REROUTES), fired + + +def test_lse_on_routes_only_variants_that_write_lse(table): + on, _ = table + served = {k: res for k, (_, row, res, _) in on.items() if row["ok"]} + assert served, "no LSE request was served" + wrong = [k for k, res in served.items() if not _writes_lse(res)] + assert not wrong, wrong + + +def test_lse_on_refusals_use_the_exact_per_path_wording(table): + """Per unified candidate: LSE on may only add the per-path LSE refusal. + A candidate that refuses with LSE off must give the same reason with it + on, or that LSE reason; one that serves with LSE off must serve with it on + unless the per-path rule says why not. The one exception is the split-KV + handover (``_HANDOVER``): there the generic tiled 2D candidate refuses + with its LDS budget and the generic split-KV candidate serves instead.""" + on, off = table + bad = [] + for key, (request, _, _, verdicts) in on.items(): + off_verdicts = off[key][3] + for candidate, (ok, why) in verdicts.items(): + off_ok, off_why = off_verdicts[candidate] + expected = _expected_lse_refusal( + candidate, request.arch, request.use_sinks, off_why + ) + if key in _HANDOVER and candidate == "attention_unified_2d": + good = off_ok and not ok and "LDS budget" in why + elif key in _HANDOVER and candidate == "attention_unified_3d": + good = not off_ok and ok + elif not off_ok: + good = not ok and why in (off_why, expected) + elif expected is None: + good = ok + else: + good = not ok and why == expected + if not good: + bad.append((key, candidate, (ok, why), expected, (off_ok, off_why))) + assert not bad, bad[:10] + + +@pytest.mark.parametrize("key", sorted(_HANDOVER)) +def test_gfx942_d256_bs32_lse_prefill_takes_split_kv(table, key): + """The tiled 2D LDS gate declines this LSE problem, so auto serves it + with split-KV and the spec's reduce writes LSE. LSE off is unchanged.""" + on, off = table + assert (off[key][1]["candidate"], off[key][1]["path"]) == ( + "attention_unified_2d", + "2d", + ) + row, result = on[key][1], on[key][2] + assert row["ok"], row + assert (row["candidate"], row["path"]) == ("attention_unified_3d", "3d") + assert result.spec.emit_lse and _writes_lse(result) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.c index de77d6d7f2bf..e1d266224f3c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.c @@ -222,6 +222,83 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) s->use_k_single_buffer = true; s->kv_cache_policy = "all"; break; + case 13: /* cfg5 + emit_lse (transposed-x8, 16 slots, slot 0 stored). */ + s->head_size = 128; + s->block_size = 32; + s->num_query_heads = 8; + s->num_kv_heads = 1; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_warps = 2; + s->block_m_per_warp = 32; + s->has_tile_size = true; + s->tile_size = 32; + s->use_mfma_32x32x8 = true; + s->use_transposed_qk_32x32 = true; + s->use_k_single_buffer = false; + s->emit_lse = true; + break; + case 14: /* emit_lse, transposed-x8 one-slot state, NQK=6. */ + s->head_size = 128; + s->block_size = 32; + s->num_query_heads = 12; + s->num_kv_heads = 2; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_warps = 2; + s->block_m_per_warp = 32; + s->has_tile_size = true; + s->tile_size = 32; + s->use_mfma_32x32x8 = true; + s->use_transposed_qk_32x32 = true; + s->use_transposed_scalar_state = true; + s->use_k_single_buffer = false; + s->emit_lse = true; + break; + case 15: /* emit_lse, non-transposed x8 Acc_lds epilogue, NQK=6. */ + s->head_size = 128; + s->block_size = 32; + s->num_query_heads = 12; + s->num_kv_heads = 2; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_warps = 2; + s->block_m_per_warp = 32; + s->has_tile_size = true; + s->tile_size = 32; + s->use_mfma_32x32x8 = true; + s->emit_lse = true; + break; + case 16: /* emit_lse, striped16 epilogue, NQK=6. */ + s->head_size = 128; + s->block_size = 32; + s->num_query_heads = 12; + s->num_kv_heads = 2; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->emit_lse = true; + break; + case 17: /* emit_lse, striped16, two atoms per warp, q-major grid. */ + s->head_size = 128; + s->block_size = 32; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->block_m_per_warp = 32; + s->use_q_major_grid = true; + s->emit_lse = true; + break; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.py index dd11dfa54cf1..fda073985652 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_emit.py @@ -9,6 +9,8 @@ # build_unified_attention_2d_tiled(spec, arch="gfx942") and prints # lower_kernel_to_llvm(kernel, arch="gfx942") to stdout so it can be # byte-compared with the C emitter gfx942_attention_tiled_2d_emit.c. +# prove_parity_binding keeps a copy of _CONFIGS in +# platform/cpp/bindings/attention_tiled_parity_configs.py; update it in step. from kernels.gfx942.attention_tiled_2d import ( UnifiedAttention2DTiledSpec, build_unified_attention_2d_tiled, @@ -212,6 +214,86 @@ use_k_single_buffer=True, kv_cache_policy="all", ), + # cfg13-cfg17: emit_lse=True, one per LSE store site. cfg13 is cfg5 + LSE + # (transposed-x8, 16 state slots, slot 0 stored). cfg14 is the transposed + # one-slot state with NQK=6 (VALID_ROWS < BLOCK_M). cfg15 is the + # non-transposed x8 Acc_lds epilogue (16 slots, lane_col32 == 0). cfg16 and + # cfg17 are the striped16 epilogue (NQK=6 guard; two atoms per warp + the + # q-major grid). + 13: dict( + head_size=128, + block_size=32, + num_query_heads=8, + num_kv_heads=1, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_warps=2, + block_m_per_warp=32, + tile_size=32, + use_mfma_32x32x8=True, + use_transposed_qk_32x32=True, + use_k_single_buffer=False, + emit_lse=True, + ), + 14: dict( + head_size=128, + block_size=32, + num_query_heads=12, + num_kv_heads=2, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_warps=2, + block_m_per_warp=32, + tile_size=32, + use_mfma_32x32x8=True, + use_transposed_qk_32x32=True, + use_transposed_scalar_state=True, + use_k_single_buffer=False, + emit_lse=True, + ), + 15: dict( + head_size=128, + block_size=32, + num_query_heads=12, + num_kv_heads=2, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_warps=2, + block_m_per_warp=32, + tile_size=32, + use_mfma_32x32x8=True, + emit_lse=True, + ), + 16: dict( + head_size=128, + block_size=32, + num_query_heads=12, + num_kv_heads=2, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + emit_lse=True, + ), + 17: dict( + head_size=128, + block_size=32, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + block_m_per_warp=32, + use_q_major_grid=True, + emit_lse=True, + ), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.c index 6101487a6520..b47641cfc848 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.c @@ -329,6 +329,48 @@ static int make_spec(int idx, rocke_unified_attention_3d_tiled_spec_t* s) s->use_decode_grid = true; s->kv_cache_policy = "stream"; break; + case 20: + /* Reduce with emit_lse (see make_reduce_spec): bf16 D128 GQA. */ + s->head_size = 128; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->num_segments = 16; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->kv_storage_dtype = NULL; + break; + case 21: + /* Reduce with emit_lse (see make_reduce_spec): fp16 D64. */ + s->head_size = 64; + s->block_size = 16; + s->num_query_heads = 8; + s->num_kv_heads = 8; + s->dtype = "fp16"; + s->num_segments = 8; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->kv_storage_dtype = NULL; + break; + case 22: + /* Decode GQA with waves_per_eu, which the reduce inherits. */ + s->head_size = 128; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->num_segments = 16; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->kv_storage_dtype = NULL; + s->use_decode_grid = true; + s->has_waves_per_eu = true; + s->waves_per_eu = 2; + break; default: return -1; } @@ -336,8 +378,10 @@ static int make_spec(int idx, rocke_unified_attention_3d_tiled_spec_t* s) } /* Derive the reduce spec from the segment spec, exactly as the Python emitter - * (head_size / num_query_heads / num_kv_heads / dtype / num_segments). */ -static void make_reduce_spec(const rocke_unified_attention_3d_tiled_spec_t* s, + * (head_size / num_query_heads / num_kv_heads / dtype / num_segments / + * waves_per_eu; configs 20-21 also set emit_lse, as _REDUCE_LSE does). */ +static void make_reduce_spec(int idx, + const rocke_unified_attention_3d_tiled_spec_t* s, rocke_unified_attention_reduce_tiled_spec_t* r) { *r = rocke_unified_attention_reduce_tiled_spec_default(); @@ -346,6 +390,9 @@ static void make_reduce_spec(const rocke_unified_attention_3d_tiled_spec_t* s, r->num_kv_heads = s->num_kv_heads; r->dtype = s->dtype; r->num_segments = s->num_segments; + r->has_waves_per_eu = s->has_waves_per_eu; + r->waves_per_eu = s->waves_per_eu; + r->emit_lse = (idx == 20 || idx == 21); } /* Emit the reduce kernel in `mode` (ll / ir / verify) after the segment. @@ -431,7 +478,7 @@ int main(int argc, char** argv) return 1; } rocke_unified_attention_reduce_tiled_spec_t r; - make_reduce_spec(&s, &r); + make_reduce_spec(idx, &s, &r); if(strcmp(mode, "ll") == 0) { diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.py index 7b57bc1443a9..5fc65d8dcfd2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_3d_emit.py @@ -17,6 +17,9 @@ # Each kernel is lowered with _native_lower(kernel, arch="gfx942") and the two # .ll texts are concatenated (segment first, then reduce) to stdout so they can # be byte-compared with the C emitter gfx942_attention_tiled_3d_emit.c. +# prove_parity_binding keeps a copy of _CONFIGS (with emit_lse set for the +# _REDUCE_LSE configs) in platform/cpp/bindings/attention_tiled_parity_configs.py; +# update it in step. import sys from kernels.gfx942.attention_tiled_3d import ( @@ -331,8 +334,52 @@ use_decode_grid=True, kv_cache_policy="stream", ), + 20: dict( + # Reduce with emit_lse (see _REDUCE_LSE): bf16 D128 GQA. + head_size=128, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + num_segments=16, + use_sinks=False, + sliding_window=0, + has_softcap=False, + kv_storage_dtype=None, + ), + 21: dict( + # Reduce with emit_lse (see _REDUCE_LSE): fp16 D64. + head_size=64, + block_size=16, + num_query_heads=8, + num_kv_heads=8, + dtype="fp16", + num_segments=8, + use_sinks=False, + sliding_window=0, + has_softcap=False, + kv_storage_dtype=None, + ), + 22: dict( + # Decode GQA with waves_per_eu, which the reduce inherits. + head_size=128, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + num_segments=16, + use_sinks=False, + sliding_window=0, + has_softcap=False, + kv_storage_dtype=None, + use_decode_grid=True, + waves_per_eu=2, + ), } +# Configs whose REDUCE kernel is built with emit_lse=True (segment unchanged). +_REDUCE_LSE = frozenset({20, 21}) + def _kernels(idx: int): if idx not in _CONFIGS: @@ -347,6 +394,8 @@ def _kernels(idx: int): num_kv_heads=cfg["num_kv_heads"], dtype=cfg["dtype"], num_segments=cfg["num_segments"], + waves_per_eu=cfg.get("waves_per_eu"), + emit_lse=idx in _REDUCE_LSE, ) red_kernel = build_unified_attention_reduce_tiled(red_spec, arch="gfx942") return seg_kernel, red_kernel diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c index 1fb78c267446..06a858db0096 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c @@ -1102,6 +1102,27 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) s->use_qblock_reverse = true; break; + /* --- idx64-67: emit_lse, one per epilogue store site, each an existing + * config plus emit_lse (lockstep with the Python _CONFIGS[64..67]). + * 64: transposed32 scalar-state (55); 65: coalesced mfma32 (32); + * 66: striped16 GQA-16 (3); 67: striped16 + ALiBi (5 + use_alibi). */ + case 64: + make_spec(55, s); + s->emit_lse = true; + break; + case 65: + make_spec(32, s); + s->emit_lse = true; + break; + case 66: + make_spec(3, s); + s->emit_lse = true; + break; + case 67: + make_spec(5, s); + s->use_alibi = true; + s->emit_lse = true; + break; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py index 5e9adcc90339..1a86134920ca 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py @@ -10,6 +10,8 @@ # build_unified_attention_2d_tiled(spec, arch="gfx950") and prints # lower_kernel_to_llvm(kernel, arch="gfx950") to stdout so it can be # byte-compared with the C emitter gfx950_attention_tiled_2d_emit.c. +# prove_parity_binding keeps a copy of _CONFIGS in +# platform/cpp/bindings/attention_tiled_parity_configs.py; update it in step. # # The config table is kept IN LOCKSTEP with the C make_spec() switch (same index # -> same UnifiedAttention2DTiledSpec). This is the "edge / feature-flag" cluster: @@ -1071,6 +1073,18 @@ ), } +# --- idx64-67: emit_lse (trailing lse_ptr + FP32 LSE store), one per epilogue +# store site, each an existing config plus emit_lse=True. +# idx64: transposed32 scalar-state (idx55 + LSE). +# idx65: coalesced mfma32 Acc_lds epilogue (idx32 + LSE). +# idx66: striped16 epilogue, GQA-16 (idx3 + LSE). +# idx67: striped16 + ALiBi (idx5 + ALiBi + LSE): the LSE row index must use +# the query head, not the hoisted ALiBi slope. +_CONFIGS[64] = {**_CONFIGS[55], "emit_lse": True} +_CONFIGS[65] = {**_CONFIGS[32], "emit_lse": True} +_CONFIGS[66] = {**_CONFIGS[3], "emit_lse": True} +_CONFIGS[67] = {**_CONFIGS[5], "use_alibi": True, "emit_lse": True} + def _spec(idx: int) -> UnifiedAttention2DTiledSpec: if idx not in _CONFIGS: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c index 7f3d4efa1579..5a3e7fadf24d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c @@ -75,6 +75,11 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) case 5: s->use_grouped_kv2_softmax = true; break; + case 6: + /* LSE output through the wrapper (trailing lse_ptr, "_lse" in the + * wrapper's kernel name). */ + s->emit_lse = true; + break; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py index d17b1cad2b35..7db60271f4d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py @@ -47,6 +47,9 @@ 3: dict(_BASE, use_transposed_half_local_pv=True, use_mfma32_skip_legacy_qreg=True), 4: dict(_BASE, use_agpr_alloc_zero=True, use_transposed_half_local_pv=True), 5: dict(_BASE, use_grouped_kv2_softmax=True), + # LSE output through the wrapper: trailing lse_ptr and the wrapper's own + # kernel name ("..._lse_fastkv_regp"). + 6: dict(_BASE, emit_lse=True), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.c index 96dc1aa70932..b8f8e4b90578 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.c @@ -303,6 +303,32 @@ static int make_spec(int idx, rocke_unified_attention_3d_tiled_spec_t* s) s->use_decode_grid = true; s->kv_cache_policy = "stream"; break; + case 18: + /* Reduce with emit_lse (see make_reduce_spec): bf16 D128 GQA. */ + s->head_size = 128; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->num_segments = 16; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->kv_storage_dtype = NULL; + break; + case 19: + /* Reduce with emit_lse (see make_reduce_spec): fp16 D64. */ + s->head_size = 64; + s->block_size = 16; + s->num_query_heads = 8; + s->num_kv_heads = 8; + s->dtype = "fp16"; + s->num_segments = 8; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->kv_storage_dtype = NULL; + break; default: return -1; } @@ -310,8 +336,10 @@ static int make_spec(int idx, rocke_unified_attention_3d_tiled_spec_t* s) } /* Derive the reduce spec from the segment spec, exactly as the Python emitter - * (head_size / num_query_heads / num_kv_heads / dtype / num_segments). */ -static void make_reduce_spec(const rocke_unified_attention_3d_tiled_spec_t* s, + * (head_size / num_query_heads / num_kv_heads / dtype / num_segments / + * waves_per_eu; configs 18-19 also set emit_lse, as _REDUCE_LSE does). */ +static void make_reduce_spec(int idx, + const rocke_unified_attention_3d_tiled_spec_t* s, rocke_unified_attention_reduce_tiled_spec_t* r) { *r = rocke_unified_attention_reduce_tiled_spec_default(); @@ -320,6 +348,9 @@ static void make_reduce_spec(const rocke_unified_attention_3d_tiled_spec_t* s, r->num_kv_heads = s->num_kv_heads; r->dtype = s->dtype; r->num_segments = s->num_segments; + r->has_waves_per_eu = s->has_waves_per_eu; + r->waves_per_eu = s->waves_per_eu; + r->emit_lse = (idx == 18 || idx == 19); } /* Build kernel via `which` (0=segment, 1=reduce), lower to .ll, print to stdout. @@ -402,7 +433,7 @@ int main(int argc, char** argv) } rocke_unified_attention_reduce_tiled_spec_t r; - make_reduce_spec(&s, &r); + make_reduce_spec(idx, &s, &r); if(strcmp(mode, "ll") == 0) { diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.py index de07cb5066db..dd7d6e6b99ab 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_3d_emit.py @@ -17,6 +17,9 @@ # Each kernel is lowered with _native_lower(kernel, arch="gfx950") and # the two .ll texts are concatenated (segment first, then reduce) to stdout so # they can be byte-compared with the C emitter gfx950_attention_tiled_3d_emit.c. +# prove_parity_binding keeps a copy of _CONFIGS (with emit_lse set for the +# _REDUCE_LSE configs) in platform/cpp/bindings/attention_tiled_parity_configs.py; +# update it in step. import sys from kernels.gfx950.attention_tiled_3d import ( @@ -297,8 +300,37 @@ use_decode_grid=True, kv_cache_policy="stream", ), + 18: dict( + # Reduce with emit_lse (see _REDUCE_LSE): bf16 D128 GQA. + head_size=128, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + num_segments=16, + use_sinks=False, + sliding_window=0, + has_softcap=False, + kv_storage_dtype=None, + ), + 19: dict( + # Reduce with emit_lse (see _REDUCE_LSE): fp16 D64. + head_size=64, + block_size=16, + num_query_heads=8, + num_kv_heads=8, + dtype="fp16", + num_segments=8, + use_sinks=False, + sliding_window=0, + has_softcap=False, + kv_storage_dtype=None, + ), } +# Configs whose REDUCE kernel is built with emit_lse=True (segment unchanged). +_REDUCE_LSE = frozenset({18, 19}) + def _emit(idx: int) -> str: if idx not in _CONFIGS: @@ -315,6 +347,8 @@ def _emit(idx: int) -> str: num_kv_heads=cfg["num_kv_heads"], dtype=cfg["dtype"], num_segments=cfg["num_segments"], + waves_per_eu=cfg.get("waves_per_eu"), + emit_lse=idx in _REDUCE_LSE, ) red_kernel = build_unified_attention_reduce_tiled(red_spec, arch="gfx950") red_ll = _native_lower(red_kernel, arch="gfx950") @@ -342,6 +376,7 @@ def main() -> int: num_kv_heads=cfg["num_kv_heads"], dtype=cfg["dtype"], num_segments=cfg["num_segments"], + emit_lse=idx in _REDUCE_LSE, ) red_kernel = build_unified_attention_reduce_tiled(red_spec, arch="gfx950") sys.stdout.write(serialize(seg_kernel)) @@ -358,6 +393,7 @@ def main() -> int: num_kv_heads=cfg["num_kv_heads"], dtype=cfg["dtype"], num_segments=cfg["num_segments"], + emit_lse=idx in _REDUCE_LSE, ) red_kernel = build_unified_attention_reduce_tiled(red_spec, arch="gfx950") sys.stdout.write("".join(str(d) + "\n" for d in verify(seg_kernel))) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_kv_cache_policy.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_kv_cache_policy.py index 4978e8bd3e64..1aedbd19ea49 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_kv_cache_policy.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_kv_cache_policy.py @@ -204,6 +204,7 @@ def test_gfx942_wide_loader_takes_all_or_stream(policy): {"use_wide_kv_load": True, "kv_cache_policy": "stream"}, {"kv_cache_policy": "nt"}, {"kv_layout": "strided", "kv_cache_policy": "stream"}, + {"kv_cache_policy": "stream", "waves_per_eu": 2}, ], ) def test_gfx942_native_builder_matches_python(kw): @@ -213,9 +214,21 @@ def test_gfx942_native_builder_matches_python(kw): engine = pytest.importorskip("rocke_engine") spec = _gfx942_spec(num_query_heads=32, num_segments=8, num_seqs=3, **kw) - # The gfx942 binding lowers the segment kernel only. + # The binding lowers the segment kernel, then the reduce it derives from the + # segment spec (waves_per_eu included). + reduce = gfx942_3d.UnifiedAttentionReduceTiledSpec( + head_size=spec.head_size, + num_query_heads=spec.num_query_heads, + num_kv_heads=spec.num_kv_heads, + dtype=spec.dtype, + num_segments=spec.num_segments, + waves_per_eu=spec.waves_per_eu, + ) expected = _lower_kernel_to_llvm_python( gfx942_3d.build_unified_attention_3d_tiled(spec, arch="gfx942"), arch="gfx942" + ) + _lower_kernel_to_llvm_python( + gfx942_3d.build_unified_attention_reduce_tiled(reduce, arch="gfx942"), + arch="gfx942", ) actual = engine.gfx942_attention_tiled_3d_lower_llvm(asdict(spec), "gfx942") assert actual == expected diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py index 203ee1d18fff..f9108bf1d9bd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py @@ -109,6 +109,7 @@ def _built_3d(problem, arch): problem.dtype, au._num_segments(problem, arch), au._select_3d_waves_per_eu(problem, arch), + problem.emit_lse, ), ) @@ -124,13 +125,14 @@ def _built_3d(problem, arch): _SEQLEN_Q = (1, 2, 64, 65, 256, 257, 512, 513, 768, 769, 1024, 1025, 1536, 2048) _SEQLEN_K = ("q", 512, 513, 1024, 1025) -# (sliding_window, use_sinks, bias) combinations; bias is one feature at a time. +# (sliding_window, use_sinks, feature) combinations; one feature at a time. _FEATURES = ( (0, False, None), (128, False, None), (0, True, None), (0, False, "softcap"), (0, False, "alibi"), + (0, False, "lse"), ) @@ -155,6 +157,7 @@ def _static_grid(arch): use_sinks=sinks, softcap=30.0 if bias == "softcap" else 0.0, use_alibi=bias == "alibi", + emit_lse=bias == "lse", ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_2d_lse_ir.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_2d_lse_ir.py new file mode 100644 index 000000000000..0eb49c39fc59 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_2d_lse_ir.py @@ -0,0 +1,341 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Host-only IR tests for opt-in LSE output on the tiled 2D unified-attention +kernels (gfx942 + gfx950). No GPU: specs are built and lowered to LLVM text.""" + +import importlib +import re + +import pytest + +from rocke.core.lower_llvm import lower_kernel_to_llvm +from rocke.core.passes import eliminate_dead_pure_ops + +_BASE = dict( + head_size=128, + block_size=64, + num_query_heads=4, + num_kv_heads=2, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + use_alibi=False, +) + +# Spec kwargs that select each epilogue branch of build_unified_attention_2d_tiled. +BRANCH_SPEC_KW = { + "gfx950": { + "transposed32": dict( + use_mfma_32x32=True, + block_m_per_warp=32, + use_transposed_qk_32x32=True, + use_transposed_scalar_state=True, + ), + "mfma32": dict(use_mfma_32x32=True, block_m_per_warp=32), + "striped16": {}, + }, + "gfx942": { + "transposed32": dict( + use_mfma_32x32x8=True, use_transposed_qk_32x32=True, block_m_per_warp=32 + ), + "mfma32": dict(use_mfma_32x32x8=True, block_m_per_warp=32), + "striped16": {}, + }, +} + +# One FP32 LSE store per (writer lane, softmax-state slot), Python-unrolled. +_EXPECTED_LSE_STORES = { + "gfx950": {"transposed32": 1, "mfma32": 16, "striped16": 4}, + "gfx942": {"transposed32": 1, "mfma32": 16, "striped16": 4}, +} + +_CASES = [ + (arch, branch) + for arch in ("gfx950", "gfx942") + for branch in ("transposed32", "mfma32", "striped16") +] +_SPEC_ARCHES = ["gfx950", "gfx942"] + + +def _module(arch): + return importlib.import_module(f"kernels.{arch}.attention_tiled_2d") + + +def _spec(arch, branch, **kw): + return _module(arch).UnifiedAttention2DTiledSpec( + **_BASE, **BRANCH_SPEC_KW[arch][branch], **kw + ) + + +def _build(spec, arch): + return _module(arch).build_unified_attention_2d_tiled(spec, arch=arch) + + +def _lower(spec, arch): + return lower_kernel_to_llvm(_build(spec, arch), arch=arch) + + +def _walk(region): + for op in region.ops: + yield op + for nested in op.regions: + yield from _walk(nested) + + +def _drop_ops(region, drop): + """Remove the ops ``drop`` selects, at any depth; return how many went.""" + kept, gone = [], 0 + for op in region.ops: + for nested in op.regions: + gone += _drop_ops(nested, drop) + if drop(op): + gone += 1 + else: + kept.append(op) + region.ops = kept + return gone + + +def _drop_dead(kernel, guards=frozenset()): + """Remove dead pure ops, and the scf.if ops in ``guards`` once empty.""" + while eliminate_dead_pure_ops(kernel.body) + _drop_ops( + kernel.body, + lambda op: id(op) in guards and not any(r.ops for r in op.regions), + ): + pass + + +def _strip_lse(kernel): + """Remove the LSE output from an LSE-on kernel, in place, on the IR before + lowering: the ``lse_ptr`` param, every global store through it, the scf.if + guards those stores leave empty, and the pure ops that only fed them. + Return how many stores went.""" + lse = next(p for p in kernel.params if p.name == "lse_ptr") + kernel.params.remove(lse) + + def is_lse_store(op): + return ( + op.name.startswith("memref.global_store") + and op.operands[0].name == f"%{lse.name}" + ) + + guards = set() + + def find_guards(region, ifs): + for op in region.ops: + if is_lse_store(op): + guards.update(ifs) + inner = ifs + (id(op),) if op.name == "scf.if" else ifs + for nested in op.regions: + find_guards(nested, inner) + + find_guards(kernel.body, ()) + stores = _drop_ops(kernel.body, is_lse_store) + _drop_dead(kernel, guards) + assert kernel.name.endswith("_lse"), kernel.name + kernel.name = kernel.name.removesuffix("_lse") + return stores + + +def _canonical_names(kernel): + """Rename every non-param value by first use, in place. Value names come + from a builder-wide counter, so LSE ops shift every later name. + + An op's result object is also the later ops' operand object, so map the + original names first and rename each object once; renaming while walking + would rekey a value on every reuse and can merge two values.""" + params = {f"%{p.name}" for p in kernel.params} + names, values = {}, {} + for op in _walk(kernel.body): + for value in (*op.operands, *op.results): + if value.name in params or id(value) in values: + continue + values[id(value)] = value + names.setdefault(value.name, f"%v{len(names)}") + for value in values.values(): + value.name = names[value.name] + + +def test_canonical_names_keep_reused_values_distinct(): + from rocke.core.ir import I32, IRBuilder + + b = IRBuilder("k") + x = b.param("x", I32) + first = b.add(x, x) + second = b.add(first, x) + third = b.mul(first, second) + # Original names already in the canonical %vN space, in the reverse order. + first.name, second.name = "%v1", "%v0" + _canonical_names(b.kernel) + assert (first.name, second.name, third.name) == ("%v0", "%v1", "%v2") + + +def _lse_store_count(ir): + """FP32 global stores. The LSE-off 2D kernels have none (bf16/fp16 O).""" + return sum( + 1 + for line in ir.splitlines() + if "store float" in line and "ptr addrspace(1)" in line + ) + + +def _define_line(ir): + return next(line for line in ir.splitlines() if line.startswith("define")) + + +@pytest.mark.parametrize("arch", _SPEC_ARCHES) +@pytest.mark.parametrize("bad", ["false", 0, 1, None]) +def test_2d_spec_rejects_non_boolean_emit_lse(arch, bad): + with pytest.raises(ValueError, match="emit_lse must be a bool"): + _spec(arch, "striped16", emit_lse=bad) + + +@pytest.mark.parametrize("arch,branch", _CASES) +def test_2d_kernel_name_gains_lse_suffix_only_when_on(arch, branch): + default = _spec(arch, branch) + off = _spec(arch, branch, emit_lse=False) + on = _spec(arch, branch, emit_lse=True) + assert default.emit_lse is False + assert off.kernel_name() == default.kernel_name() + assert not default.kernel_name().endswith("_lse") + assert on.kernel_name() == default.kernel_name() + "_lse" + + +# gfx942 also builds branch A with one softmax-state slot (production flash +# specs set it) and the q-major grid; the name rule must hold for those too. +@pytest.mark.parametrize( + "branch,extra", + [ + ("transposed32", {}), + ("transposed32", {"use_transposed_scalar_state": True}), + ("mfma32", {}), + ("striped16", {}), + ("striped16", {"use_q_major_grid": True}), + ], +) +def test_gfx942_lse_name_is_off_name_plus_lse_suffix(branch, extra): + off = _spec("gfx942", branch, **extra) + on = _spec("gfx942", branch, emit_lse=True, **extra) + assert off.emit_lse is False + assert not off.kernel_name().endswith("_lse") + assert on.kernel_name() == off.kernel_name() + "_lse" + + +@pytest.mark.parametrize("arch,branch", _CASES) +def test_2d_lse_off_ir_has_no_lse_param_or_store(arch, branch): + default_ir = _lower(_spec(arch, branch), arch) + off_ir = _lower(_spec(arch, branch, emit_lse=False), arch) + assert off_ir == default_ir + assert "lse_ptr" not in off_ir + assert _lse_store_count(off_ir) == 0 + + +@pytest.mark.parametrize("arch,branch", _CASES) +def test_2d_lse_on_ir_minus_lse_output_is_lse_off_ir(arch, branch): + """LSE on adds only the LSE output: with its param, stores, their guards + and the pure ops that fed only them removed, it lowers to the LSE-off + text. Dead pure ops go from both sides (LLVM drops them anyway). The + LSE-off text itself is pinned by rocke_representative_ir_sha256.json.""" + off = _build(_spec(arch, branch), arch) + on = _build(_spec(arch, branch, emit_lse=True), arch) + assert _strip_lse(on) == _EXPECTED_LSE_STORES[arch][branch] + _drop_dead(off) + for kernel in (off, on): + _canonical_names(kernel) + off_ir = lower_kernel_to_llvm(off, arch=arch) + on_ir = lower_kernel_to_llvm(on, arch=arch) + assert on_ir == off_ir, f"LSE on changes more than its output: {arch}/{branch}" + + +@pytest.mark.parametrize("arch,branch", _CASES) +def test_2d_lse_on_appends_lse_ptr_last_and_stores_once_per_slot(arch, branch): + on_ir = _lower(_spec(arch, branch, emit_lse=True), arch) + define = _define_line(on_ir) + assert re.search( + r"ptr addrspace\(1\) noalias writeonly nocapture align 4 %lse_ptr\) #\d+ \{$", + define, + ), define + assert _lse_store_count(on_ir) == _EXPECTED_LSE_STORES[arch][branch] + # natural_lse_from_log2_stats: (m + log2 l) * ln2 + assert "@llvm.log2.f32" in on_ir + + +# Additive-bias params are appended after qq_bias_stride_0; branch B rejects bias. +@pytest.mark.parametrize( + "arch,branch", [("gfx950", "transposed32"), ("gfx950", "striped16")] +) +def test_2d_lse_ptr_stays_last_with_additive_bias(arch, branch): + define = _define_line( + _lower(_spec(arch, branch, emit_lse=True, use_additive_bias=True), arch) + ) + assert "%additive_bias_sq_stride, ptr addrspace(1)" in define, define + assert re.search(r"%lse_ptr\) #\d+ \{$", define), define + + +@pytest.mark.parametrize( + "branch,extra,stores", + [ + # Production flash specs set use_transposed_scalar_state (1 state slot). + ("transposed32", {"use_transposed_scalar_state": True}, 1), + # Swapped launch grid: same kv_head_idx / q_block names, same index. + ("striped16", {"use_q_major_grid": True}, 4), + ("transposed32", {"use_q_major_grid": True}, 1), + # Two 16-row atoms per warp: 8 state slots per lane. + ("striped16", {"block_m_per_warp": 32}, 8), + ], +) +def test_gfx942_lse_variants_store_once_per_row_slot(branch, extra, stores): + off_ir = _lower(_spec("gfx942", branch, **extra), "gfx942") + on_ir = _lower(_spec("gfx942", branch, emit_lse=True, **extra), "gfx942") + assert "lse_ptr" not in off_ir + assert _lse_store_count(off_ir) == 0 + assert _lse_store_count(on_ir) == stores + assert "@llvm.log2.f32" in on_ir + + +def _gfx942_branch_of(spec): + if spec.use_mfma_32x32x8: + return "transposed32" if spec.use_transposed_qk_32x32 else "mfma32" + return "striped16" + + +@pytest.mark.parametrize( + "fields,branch", + [ + # Mandate shape, bf16 flash (_spec_gfx942_bf16_flash). + ( + dict(dtype="bf16", total_q=8192, num_seqs=1, max_seqlen_q=8192), + "transposed32", + ), + # fp16 flash (_spec_gfx942_fp16_flash). + ( + dict(dtype="fp16", total_q=8192, num_seqs=1, max_seqlen_q=8192), + "transposed32", + ), + # Decode-like 2D shape: generic narrow (_spec_generic_2d_non_gfx950). + (dict(dtype="bf16", total_q=4, num_seqs=4, max_seqlen_q=1), "striped16"), + # D256 bf16 prefill: generic narrow when the 4-warp cohort is not used. + ( + dict( + dtype="bf16", head_size=256, total_q=8192, num_seqs=1, max_seqlen_q=8192 + ), + "striped16", + ), + ], +) +def test_gfx942_production_specs_use_branches_the_numeric_gate_covers(fields, branch): + """Production gfx942 2D specs land on branch A or C, the branches the GPU + numeric gate runs by default. Branch B (non-transposed x8) is reachable + only through explicit or tuning specs, which the gate pins separately.""" + from kernels.common import attention_unified as au + + base = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_k=8192, + ) + problem = au.UnifiedAttentionProblem(**{**base, **fields}) + assert _gfx942_branch_of(au._tiled_spec_from_problem(problem, "gfx942")) == branch diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse.py index 582bc1a50c08..327208511a78 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse.py @@ -1,7 +1,7 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Host-only contract for the optional unified split-KV (3D) LSE output.""" +"""Host-only contract for the optional unified (split-KV and tiled 2D) LSE output.""" import importlib import sys @@ -76,6 +76,22 @@ def _problem(**overrides): return au.UnifiedAttentionProblem(**fields) +def _prefill_problem(arch, **overrides): + """Single-sequence causal prefill (Sq = Sk = 256, D128 bf16, GQA 32/8). + + ``select_path()`` routes it to tiled 2D, and the ``arch`` tiled gate serves + it with LSE off, so any LSE refusal seen in a test is the LSE rule's.""" + problem = _problem( + total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, **overrides + ) + assert problem.select_path() == "2d" + ok, why = au.supports_native_unified_attention_tiled( + replace(problem, emit_lse=False), arch + ) + assert ok, why + return problem + + @pytest.mark.parametrize("bad", ["false", 0, 1, None]) def test_problem_rejects_non_boolean_emit_lse(bad): with pytest.raises(ValueError, match="emit_lse must be a bool"): @@ -91,9 +107,77 @@ def test_split_kv_gate_admits_lse(arch): @pytest.mark.parametrize("arch", _ARCHES) -def test_tiled_2d_gate_declines_lse(arch): - ok, why = au.supports_native_unified_attention_tiled(_problem(emit_lse=True), arch) - assert not ok and "LSE" in why +def test_tiled_2d_gate_admits_lse_for_prefill(arch): + ok, why = au.supports_native_unified_attention_tiled( + _prefill_problem(arch, emit_lse=True), arch + ) + assert ok, why + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_tiled_2d_gate_declines_lse_with_sinks(arch): + ok, why = au.supports_native_unified_attention_tiled( + _prefill_problem(arch, emit_lse=True, use_sinks=True), arch + ) + assert not ok + assert ( + "LSE output with attention sinks is not implemented on unified attention" in why + ) + + +def test_tiled_2d_gate_declines_lse_off_cdna3_cdna4(): + ok, why = au.supports_native_unified_attention_tiled( + _problem( + total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, emit_lse=True + ), + "gfx1250", + ) + assert not ok and "LSE output is not implemented on gfx1250" in why + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_scalar_gate_refuses_lse(arch): + """The scalar kernel writes no LSE: it must never be the auto tail for one.""" + problem = _prefill_problem(arch) + ok, why = au.supports_native_unified_attention(problem, arch) + assert ok, why + ok, why = au.supports_native_unified_attention( + replace(problem, emit_lse=True), arch + ) + assert not ok and "scalar 2D" in why + + +def test_gfx942_4warp_cohort_is_ineligible_when_lse_is_on(): + """D256 bf16 causal prefill is the 4-warp GQA cohort with LSE off. With LSE + on it falls to the default gfx942 builder, which writes LSE.""" + d256 = _problem( + total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, head_size=256 + ) + assert au._gfx942_4warp_route(d256, "gfx942") is not None + assert au._gfx942_4warp_route(replace(d256, emit_lse=True), "gfx942") is None + swa = _problem( + total_q=256, + num_seqs=1, + max_seqlen_q=256, + max_seqlen_k=256, + sliding_window=64, + ) + assert au._gfx942_4warp_route(swa, "gfx942") is not None + assert au._gfx942_4warp_route(replace(swa, emit_lse=True), "gfx942") is None + + +def test_gfx942_4warp_builder_refuses_an_lse_spec(): + """The last gate: a spec built directly (no runner, no dispatcher) with + emit_lse=True must not become a 4-warp kernel that has no LSE store.""" + from kernels.gfx942.attention_tiled_2d import build_gfx942_4warp_gqa + + d256 = _problem( + total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, head_size=256 + ) + spec = au._tiled_spec_from_problem(d256, "gfx942") + assert "lse" not in build_gfx942_4warp_gqa(spec, arch="gfx942").name + with pytest.raises(NotImplementedError, match="4-warp GQA builder"): + build_gfx942_4warp_gqa(replace(spec, emit_lse=True), arch="gfx942") def test_split_kv_gate_declines_lse_off_cdna3_cdna4(): @@ -127,9 +211,13 @@ def test_allow_unsupported_cannot_bypass_lse_eligibility(arch): ) assert not ok and "sink" in why ok, why = au._explicit_path_supported( - _problem(emit_lse=True), _BypassSpec(), "2d", arch + _problem(emit_lse=True, use_sinks=True), _BypassSpec(), "2d", arch ) - assert not ok and "LSE" in why + assert not ok and "sink" in why + ok, why = au._explicit_path_supported( + _problem(emit_lse=True), _BypassSpec(), "2d", "gfx1250" + ) + assert not ok and "gfx1250" in why # The bypass still works for everything else. ok, _ = au._explicit_path_supported( _problem(use_sinks=True), _BypassSpec(), "3d", arch @@ -150,6 +238,92 @@ def test_reduce_signature_appends_lse_last(): assert on == off + ["lse_ptr"] +@pytest.mark.parametrize("arch", _ARCHES) +def test_tiled_spec_carries_lse_only_when_requested(arch): + off = au._tiled_spec_from_problem(_prefill_problem(arch), arch) + on = au._tiled_spec_from_problem(_prefill_problem(arch, emit_lse=True), arch) + assert off.emit_lse is False and on.emit_lse is True + assert replace(on, emit_lse=False) == off + assert on.kernel_name() == off.kernel_name() + "_lse" + + +@pytest.mark.parametrize("arch", _ARCHES) +def test_tiled_cache_key_keeps_lse_off_and_suffixes_lse_on(arch): + off = au._tiled_cache_key(_prefill_problem(arch), arch) + on = au._tiled_cache_key(_prefill_problem(arch, emit_lse=True), arch) + assert off[-1] != "lse" + assert on == off + ("lse",) + + +def test_tiled_cache_key_rekeys_the_gfx942_4warp_cohort_under_lse(): + """D256 bf16 prefill is a 4-warp cohort with LSE off; with LSE on the + default builder serves it, so resolved entries change, not just the suffix.""" + d256 = _problem( + total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, head_size=256 + ) + off = au._tiled_cache_key(d256, "gfx942") + on = au._tiled_cache_key(replace(d256, emit_lse=True), "gfx942") + assert on[-1] == "lse" and on[:-1] != off + + +def test_attn_signature_appends_lse_last(): + kw = dict(include_bt_stride=True, include_qq_bias_stride=True) + off = [p["name"] for p in au._attn_signature("bf16", **kw)] + on = [p["name"] for p in au._attn_signature("bf16", emit_lse=True, **kw)] + assert on == off + ["lse_ptr"] + assert off[-1] == "qq_bias_stride_0" + # The additive-bias params come after qq_bias_stride_0; lse_ptr + # stays last, after them, matching the gfx950 kernel's declaration order. + kw["include_additive_bias"] = True + off = [p["name"] for p in au._attn_signature("bf16", **kw)] + on = [p["name"] for p in au._attn_signature("bf16", emit_lse=True, **kw)] + assert on == off + ["lse_ptr"] + assert off[-1] == "additive_bias_sq_stride" + + +def _values(problem, **extra): + return au._attn_values( + problem=problem, + q=1, + k=2, + v=3, + out=4, + cu_seqlens_q=5, + seqused_k=6, + softmax_scale=0.125, + block_table=7, + softcap=0.0, + sinks=0, + bt_stride=16, + include_bt_stride=True, + include_qq_bias_stride=True, + kv_addr_limit=None, + **extra, + ) + + +def test_attn_values_bind_lse_only_when_requested(): + lse = object() + assert "lse_ptr" not in _values(_prefill_problem("gfx950")) + on = _values(_prefill_problem("gfx950", emit_lse=True), lse=lse) + assert on["lse_ptr"] is lse + assert list(on)[-1] == "lse_ptr" + + +@pytest.mark.parametrize("emit_lse", [False, True]) +def test_2d_launcher_signature_follows_the_problem(monkeypatch, emit_lse): + """A cached HSACO stands in for the compile; only the launcher is built.""" + problem = _prefill_problem("gfx942", emit_lse=emit_lse) + key = au._tiled_cache_key(problem, "gfx942") + monkeypatch.setattr(au, "_2D_LAUNCHERS", {}) + monkeypatch.setattr(au, "_ATTN_TILED_CACHE", {key: (b"", "kernel")}) + monkeypatch.setattr(au, "KernelLauncher", lambda **kw: SimpleNamespace(**kw)) + launcher = au._get_2d_launcher(problem, key, arch="gfx942") + names = [p["name"] for p in launcher.signature] + assert (names[-1] == "lse_ptr") is emit_lse + assert names.count("lse_ptr") == int(emit_lse) + + def test_validator_contract_without_gpu(): with pytest.raises(ValueError, match="caller-owned lse"): au.validate_unified_lse( @@ -187,7 +361,7 @@ def test_validator_contract_without_gpu(): @pytest.fixture def runner(monkeypatch): - state = SimpleNamespace(red_vals=None) + state = SimpleNamespace(red_vals=None, vals2d=None, graphed=0, validated=[]) def empty(shape, *, dtype, device): return SimpleNamespace(shape=shape, device=device) @@ -214,13 +388,57 @@ def launch(prepared, seg_vals, red_vals, stream, *, capturing): monkeypatch.setattr(au, "_launch_3d_pipeline", launch) - # Any 2D or scalar launch for an LSE request is a silent output drop. + # Record what the shared LSE buffer validator was given (it stays a stub: + # the routing contract is checked without a GPU). + monkeypatch.setattr( + au, "validate_unified_lse", lambda *a, **k: state.validated.append(k) + ) + + # The scalar kernel writes no LSE: any scalar launch here is an output drop. def never(*a, **k): - raise AssertionError("an LSE request reached a path that writes no LSE") + raise AssertionError("a request reached the scalar path, which writes no LSE") + + monkeypatch.setattr(au, "_get_scalar_launcher", never) + + # The tiled 2D launch records the bound arguments and the signature it was + # built with. + def get_2d_launcher(problem, key, *, arch, tuning_spec=None): + sig = [ + p["name"] + for p in au._attn_signature( + problem.dtype, + include_bt_stride=True, + include_qq_bias_stride=True, + include_additive_bias=problem.use_additive_bias, + emit_lse=problem.emit_lse, + ) + ] + state.sig2d, state.key2d = sig, key + + def launch(vals, *, config): + state.vals2d = dict(vals) + + return launch + + monkeypatch.setattr(au, "_get_2d_launcher", get_2d_launcher) + # develop's tiled launch passes meta.kv_addr_limit to _attn_values (A:6333). + monkeypatch.setattr( + au, + "_get_2d_launch_meta", + lambda *a, **k: SimpleNamespace( + grid=(1, 1, 1), block=(64, 1, 1), kv_addr_limit=None + ), + ) - for name in ("_run_2d_graphed", "_get_2d_launcher", "_get_scalar_launcher"): - monkeypatch.setattr(au, name, never) - # Make the 2D graph fast path eligible, so a refusal placed after it fails. + # The gfx942 2D graph replays a cache hit before any gate, and its key does + # not encode emit_lse: an LSE request must never reach it. + def graphed(problem, **kw): + assert not problem.emit_lse, "an LSE request reached the 2D graph" + state.graphed += 1 + return "graphed" + + monkeypatch.setattr(au, "_run_2d_graphed", graphed) + # Make the 2D graph fast path eligible, so a missing bypass fails. monkeypatch.setattr(au, "_enable_2d_graph_replay", lambda problem, arch: True) return state @@ -254,36 +472,36 @@ def test_runner_without_lse_launches_the_unchanged_reduce(runner): @pytest.mark.parametrize("backend", ["auto", "tiled"]) -def test_runner_refuses_lse_on_a_2d_routed_shape_before_any_2d_path(runner, backend): - prefill = _problem(total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256) - assert prefill.select_path() == "2d" - with pytest.raises(NotImplementedError, match="LSE"): - _run(replace(prefill, emit_lse=True), lse=object(), backend=backend) - assert runner.red_vals is None +def test_runner_launches_lse_on_a_2d_routed_shape_without_the_graph(runner, backend): + lse = object() + _run(_prefill_problem("gfx942", emit_lse=True), lse=lse, backend=backend) + assert runner.vals2d["lse_ptr"] is lse + assert runner.sig2d[-1] == "lse_ptr" + assert runner.key2d[-1] == "lse" + assert runner.graphed == 0 and runner.red_vals is None -@pytest.mark.parametrize( - "top_left,hinted", [(False, True), (True, False)], ids=["servable", "mask"] -) -def test_auto_refusal_names_backend_3d_only_when_split_kv_serves( - runner, top_left, hinted -): - """``auto`` routes this prefill to 2D, which writes no LSE. The refusal - points at backend='3d' only when split-KV accepts the problem: gfx942 - split-KV declines the top-left mask, so there the hint would be wrong.""" +def test_runner_without_lse_keeps_the_2d_graph(runner): + assert _run(_prefill_problem("gfx942"), backend="auto") == "graphed" + assert runner.graphed == 1 and runner.vals2d is None + + +def test_auto_lse_refusal_names_why_both_lse_paths_declined(runner): + """gfx942 tiled 2D and split-KV both decline the top-left mask; the scalar + tail writes no LSE, so the refusal names both reasons and nothing launches.""" prefill = _problem( total_q=256, num_seqs=1, max_seqlen_q=256, max_seqlen_k=256, emit_lse=True, - causal_top_left=top_left, + causal_top_left=True, ) assert prefill.select_path() == "2d" with pytest.raises(NotImplementedError, match="LSE") as refused: _run(prefill, lse=object()) - assert ("backend='3d'" in str(refused.value)) is hinted - assert runner.red_vals is None + assert "tiled 2D:" in str(refused.value) + assert runner.red_vals is None and runner.vals2d is None def test_runner_reports_why_split_kv_declined_an_lse_request(runner): @@ -293,10 +511,51 @@ def test_runner_reports_why_split_kv_declined_an_lse_request(runner): assert top_left.select_path() == "3d" with pytest.raises(NotImplementedError, match="split-KV.*top-left") as refused: _run(top_left, lse=object()) - assert "2D" not in str(refused.value) + assert "tiled 2D:" in str(refused.value) + assert runner.vals2d is None + assert runner.red_vals is None + + +def _gfx942_d256_bs32_prefill(**overrides): + """gfx942 D256 bf16 prefill at block size 32: ``select_path()`` is 2D. + LSE off, the 4-warp builder serves it; LSE on, the default tiled 2D spec + exceeds the LDS budget.""" + problem = _problem( + total_q=2048, + num_seqs=1, + head_size=256, + block_size=32, + max_seqlen_q=2048, + max_seqlen_k=2048, + **overrides, + ) + assert problem.select_path() == "2d" + return problem + + +def test_auto_hands_an_lse_request_tiled_2d_declines_to_split_kv(runner): + problem = _gfx942_d256_bs32_prefill(emit_lse=True) + ok, why = au.supports_native_unified_attention_tiled(problem, "gfx942") + assert not ok and "LDS budget" in why + lse = object() + _run(problem, lse=lse) + assert runner.red_vals["lse_ptr"] is lse + assert runner.vals2d is None and runner.graphed == 0 + + +def test_auto_without_lse_keeps_tiled_2d_for_that_shape(runner): + assert _run(_gfx942_d256_bs32_prefill(), backend="auto") == "graphed" assert runner.red_vals is None +def test_tiled_backend_still_refuses_an_lse_request_its_gate_declines(runner): + """Only ``auto`` may hand the request to split-KV; ``tiled`` was asked for + the 2D kernel and refuses with its gate's reason.""" + with pytest.raises(NotImplementedError, match="LDS budget"): + _run(_gfx942_d256_bs32_prefill(emit_lse=True), lse=object(), backend="tiled") + assert runner.red_vals is None and runner.vals2d is None + + def test_runner_refuses_lse_for_strided_kv(runner): with pytest.raises(NotImplementedError, match="strided"): _run(_problem(emit_lse=True), lse=object(), kv_layout="strided") @@ -337,14 +596,54 @@ def test_runner_rejects_a_tuning_spec_that_disagrees_on_lse( assert runner.red_vals is None -def test_runner_refuses_lse_for_an_explicit_2d_tuning_spec(runner): - """A 2D spec cannot write LSE: refused as NotImplementedError, the same - error the backend paths raise, not a spec/problem-disagreement ValueError.""" +def _tuning_spec_2d(*, kernel_emit_lse, builder_kind="tiled"): spec = _tuning_spec(reduce_emit_lse=False) - spec.path, spec.reduce_spec = "2d", None - with pytest.raises(NotImplementedError, match="LSE"): - _run(_problem(emit_lse=True), lse=object(), tuning_spec=spec) - assert runner.red_vals is None + spec.path, spec.reduce_spec, spec.builder_kind = "2d", None, builder_kind + spec.kernel_spec = SimpleNamespace(emit_lse=kernel_emit_lse, kv_storage_dtype=None) + return spec + + +@pytest.mark.parametrize( + "problem_lse,spec_lse", [(True, False), (False, True)], ids=["drop", "extra"] +) +def test_runner_rejects_a_2d_tuning_spec_that_disagrees_on_lse( + runner, problem_lse, spec_lse +): + with pytest.raises(ValueError, match="disagrees with problem.emit_lse"): + _run( + _prefill_problem("gfx942", emit_lse=problem_lse), + lse=object() if problem_lse else None, + tuning_spec=_tuning_spec_2d(kernel_emit_lse=spec_lse), + ) + assert runner.vals2d is None + + +def test_runner_launches_an_agreeing_2d_tuning_spec_with_lse(runner): + lse = object() + _run( + _prefill_problem("gfx942", emit_lse=True), + lse=lse, + tuning_spec=_tuning_spec_2d(kernel_emit_lse=True), + ) + assert runner.vals2d["lse_ptr"] is lse and runner.graphed == 0 + + +def test_runner_refuses_lse_for_a_gfx942_4warp_tuning_spec(runner): + spec = _tuning_spec_2d(kernel_emit_lse=True, builder_kind="gfx942_4warp_gqa") + with pytest.raises(NotImplementedError, match="4-warp"): + _run(_prefill_problem("gfx942", emit_lse=True), lse=object(), tuning_spec=spec) + assert runner.vals2d is None + + +def test_runner_validates_the_2d_lse_buffer_with_the_shared_overlap_set(runner): + """The LSE buffer must not overlap the additive bias either. The + runner calls validate_unified_lse (_check_lse_buffer) once, before any + path is chosen, so the tiled 2D path reuses it rather than adding its own.""" + lse = object() + _run(_prefill_problem("gfx942", emit_lse=True), lse=lse) + (checked,) = runner.validated + assert "additive_bias" in checked + assert runner.vals2d["lse_ptr"] is lse def test_runner_rejects_a_strided_tuning_spec_whose_reduce_emits_lse(runner): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse_numeric.py index 9db2042f6194..a04aab94b16a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_unified_lse_numeric.py @@ -1,11 +1,13 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""Unified split-KV (3D) LSE output vs an independent FP32 paged reference.""" +"""Unified paged attention LSE output -- split-KV (3D) reduce and tiled 2D +epilogues on gfx942/gfx950 -- vs an independent FP32 paged reference.""" import dataclasses import itertools import math +from unittest import mock import pytest @@ -15,6 +17,7 @@ UnifiedAttentionProblem, run_unified_attention_torch, ) +from kernels.common import attention_unified as au # noqa: E402 def _arch(): @@ -34,7 +37,16 @@ def _tolerance(dtype): def _paged_reference( - q, k_cache, v_cache, q_lens, kv_lens, block_table, scale, softcap, sliding_window=0 + q, + k_cache, + v_cache, + q_lens, + kv_lens, + block_table, + scale, + softcap, + sliding_window=0, + additive_bias=None, ): """FP32 O and natural-log LSE with the unified bottom-right causal mask (key k visible to query p iff k <= diag, diag = p + kv_len - q_len) and, when @@ -54,6 +66,9 @@ def _paged_reference( s = torch.matmul(qi, k.transpose(-1, -2)) * scale if softcap > 0: s = softcap * torch.tanh(s / softcap) + if additive_bias is not None: + # [seq, q_head, q_row, key], added after softcap, as the kernels do. + s = s + additive_bias[i, :, :ql, :kl] p = torch.arange(ql, device=q.device).view(-1, 1) kk = torch.arange(kl, device=q.device).view(1, -1) diag = p + (kl - ql) @@ -119,7 +134,20 @@ def _case( return problem, q, k_cache, v_cache, cu, seqused, block_table -def _run(problem, q, k_cache, v_cache, cu, seqused, block_table, *, lse, softcap): +def _run( + problem, + q, + k_cache, + v_cache, + cu, + seqused, + block_table, + *, + lse, + softcap, + backend="3d", + tuning_spec=None, +): out = torch.empty_like(q) run_unified_attention_torch( problem=problem, @@ -132,7 +160,8 @@ def _run(problem, q, k_cache, v_cache, cu, seqused, block_table, *, lse, softcap softmax_scale=1.0 / math.sqrt(q.shape[-1]), block_table=block_table, softcap=softcap, - backend="3d", + backend=backend, + tuning_spec=tuning_spec, lse=lse, ) torch.cuda.synchronize() @@ -183,6 +212,39 @@ def test_split_kv_output_and_lse(dtype, d, q_lens, kv_lens): assert (lse[finite] - lse_ref[finite]).abs().max().item() < tol +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +def test_split_kv_lse_nan_query_row_stays_nan(dtype): + """The split-KV reduce follows the tiled 2D rule: a NaN query row gives + NaN O and NaN LSE for that (token, head) only, never the keyless -inf. + Its merged sum is NaN, and !(l == 0) keeps row_has_mass true.""" + q_lens, kv_lens = [1, 1, 1, 1], [4096, 777, 1500, 2049] + nan_token, nan_head = 1, 5 + problem, q, kc, vc, cu, su, bt = _case(q_lens, kv_lens, dtype=dtype) + q[nan_token, nan_head] = float("nan") + lse = torch.full( + (problem.total_q, problem.num_query_heads, 1), float("inf"), device="cuda" + ) + out = _run(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(q.shape[-1]), 0.0 + ) + assert not torch.isposinf(lse).any(), "an LSE row was left unwritten" + assert torch.isnan(lse[nan_token, nan_head]).all(), "NaN row got a non-NaN LSE" + assert torch.isnan(out[nan_token, nan_head]).all(), "NaN row was hidden in O" + rest = torch.ones( + problem.total_q, problem.num_query_heads, dtype=torch.bool, device="cuda" + ) + rest[nan_token, nan_head] = False + assert torch.isfinite(lse.squeeze(-1)[rest]).all(), "NaN leaked into another LSE" + assert torch.isfinite(out[rest]).all(), "NaN leaked into another (row, head)" + assert (out[rest].float() - o_ref[rest]).abs().max().item() < _tolerance(dtype) + torch.testing.assert_close( + lse.squeeze(-1)[rest], lse_ref.squeeze(-1)[rest], atol=_LSE_ATOL, rtol=_LSE_RTOL + ) + + @requires_cdna3_or_4 @pytest.mark.gpu def test_split_kv_lse_with_softcap(): @@ -211,6 +273,171 @@ def _nan_lse(problem): ) +# --------------------------------------------------------------------------- +# Tiled 2D LSE (gfx942 / gfx950 build_unified_attention_2d_tiled) +# --------------------------------------------------------------------------- + +_LSE_ATOL = 1e-3 +_LSE_RTOL = 1e-3 + +# Prefill shapes the tiled 2D path serves. With 32 query / 8 KV heads every one +# routes to 2D under backend="auto" (UnifiedAttentionProblem.select_path). +_PREFILL_SHAPES = [ + # (q_lens, kv_lens) what it covers + ([256], [256]), # s_q == s_kv; the gfx942 2D-graph-eligible shape + ([37, 300, 1], [37, 812, 900]), # ragged, causal s_q < s_kv, partial pages + ( + [5], + [3], + ), # bottom-right s_q > s_kv: rows 0,1 see no key (m=-inf, l=0 after Part 0) + ([3, 64], [0, 64]), # kv_len == 0: zero-trip KV loop (m=-inf, l=1) +] + + +def _forbid_split_kv(*_args, **_kwargs): + raise AssertionError("a tiled 2D LSE test reached the split-KV (3D) path") + + +def _case2d(q_lens, kv_lens, **kw): + """``_case`` for a tiled 2D LSE launch; fails fast, with the gate's reason, + if the tiled 2D gate refuses the LSE problem on this GPU.""" + case = _case(q_lens, kv_lens, **kw) + ok, why = au.supports_native_unified_attention_tiled(case[0], _arch()) + assert ok, f"the tiled 2D gate refuses this LSE problem: {why}" + return case + + +def _run2d( + problem, + q, + k_cache, + v_cache, + cu, + seqused, + block_table, + *, + lse, + softcap, + backend, + tuning_spec=None, +): + """One 2D launch. The split-KV runner is patched to fail, so a pass proves + the tiled 2D path wrote O (and LSE).""" + with mock.patch.object(au, "_run_3d_tiled", _forbid_split_kv): + return _run( + problem, + q, + k_cache, + v_cache, + cu, + seqused, + block_table, + lse=lse, + softcap=softcap, + backend=backend, + tuning_spec=tuning_spec, + ) + + +def _run2d_lse_off( + problem, q, k_cache, v_cache, cu, seqused, block_table, *, softcap, backend +): + """The LSE-off launch of the kernel the LSE-on request runs. With LSE on, + the gfx942 4-warp cohorts fall back to the default builder and the gfx942 + 2D graph is bypassed; force both here so ``torch.equal`` compares + one kernel with and without its LSE store. Both are no-ops on gfx950.""" + off = dataclasses.replace(problem, emit_lse=False) + with mock.patch.object( + au, "_gfx942_4warp_eligible", lambda problem, arch: False + ), mock.patch.object(au, "_enable_2d_graph_replay", lambda problem, arch: False): + return _run2d( + off, + q, + k_cache, + v_cache, + cu, + seqused, + block_table, + lse=None, + softcap=softcap, + backend=backend, + ) + + +def _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype): + assert not torch.isnan(lse).any(), "an LSE row was left unwritten" + assert torch.equal(out, off), "LSE-on changed O" + assert torch.isfinite(out).all(), "non-finite O" + assert (out.float() - o_ref).abs().max().item() < _tolerance(dtype) + keyless = torch.isneginf(lse_ref.squeeze(-1)) # [total_q, Hq] + assert (out[keyless] == 0).all(), "a row with no visible key has O != 0" + assert torch.equal( + torch.isneginf(lse), torch.isneginf(lse_ref) + ), "the -inf (keyless) LSE rows differ from the reference" + assert not torch.isposinf(lse).any(), "+inf LSE" + finite = torch.isfinite(lse_ref) + torch.testing.assert_close( + lse[finite], lse_ref[finite], atol=_LSE_ATOL, rtol=_LSE_RTOL + ) + + +# One explicit dispatcher geometry per tiled 2D epilogue branch (flags = +# dispatch/attention/axes.py _CODEPATH_KNOBS). +_BRANCH_VARIANT = { + "gfx942": { + "transposed32": "transposed_x8_nw1_mw32_t2xb_llvm", + "mfma32": "wide32x8_nw1_mw32_t2xb_llvm", + "striped16": "narrow_nw1_mw16_t1xb_llvm", + }, + "gfx950": { + "transposed32": "transposed32_nw1_mw32_t2xb_llvm", + "mfma32": "wide32_nw1_mw32_t2xb_llvm", + "striped16": "narrow_nw1_mw16_t1xb_llvm", + }, +} +# A: transposed 32x32 (16 broadcast slots, and the 1-slot scalar state that +# auto routing ships); B: non-transposed 32x32 Acc_lds; C: striped 16x16. +_BRANCHES = ("transposed32", "transposed32_s1", "mfma32", "striped16") + + +def _epilogue_branch(kernel_spec) -> str: + wide = bool( + getattr(kernel_spec, "use_mfma_32x32", False) + or getattr(kernel_spec, "use_mfma_32x32x8", False) + ) + if not wide: + return "striped16" + return "transposed32" if kernel_spec.use_transposed_qk_32x32 else "mfma32" + + +def _branch_spec(problem, branch): + """The explicit tuning spec that forces ``branch`` for ``problem``.""" + from dispatch.attention.unified_rules import canonicalize_tuning_spec + + arch = _arch() + if arch == "gfx942": + from dispatch.attention.gfx942_unified import ( + GFX942_TUNING_VARIANTS as variants, + ) + else: + from dispatch.attention.gfx950_unified import ( + GFX950_TUNING_VARIANTS as variants, + ) + base = branch.removesuffix("_s1") + spec_id = f"{arch}_u2d_{_BRANCH_VARIANT[arch][base]}" + variant = next(v for v in variants if v.spec_id == spec_id) + knobs = {"use_transposed_scalar_state": True} if branch.endswith("_s1") else {} + spec, why = canonicalize_tuning_spec(problem, variant, knobs) + assert spec is not None, f"{spec_id} {knobs}: {why}" + ks = spec.kernel_spec + assert _epilogue_branch(ks) == base, f"{spec_id} selects {_epilogue_branch(ks)}" + assert bool(getattr(ks, "use_transposed_scalar_state", False)) == ( + branch.endswith("_s1") + ) + assert ks.emit_lse is problem.emit_lse + return spec + + # Decode with ragged KV and spec-decode: the shapes the split-KV path serves. _DECODE_SHAPES = [([1, 1, 1, 1], [4096, 777, 1500, 2049]), ([4, 2], [3000, 1029])] @@ -328,40 +555,75 @@ def test_masked_segment_does_not_underflow_the_merge(q_len, kv_len, emit_lse): @requires_cdna3_or_4 @pytest.mark.gpu -def test_auto_backend_refuses_lse_on_a_2d_routed_shape(): - problem, q, kc, vc, cu, su, bt = _case([256], [256]) - lse = torch.empty((problem.total_q, problem.num_query_heads, 1), device="cuda") - out = torch.empty_like(q) - with pytest.raises(NotImplementedError, match="LSE"): - run_unified_attention_torch( - problem=problem, - q=q, - k=kc, - v=vc, - out=out, - cu_seqlens_q=cu, - seqused_k=su, - softmax_scale=1.0 / math.sqrt(128), - block_table=bt, - softcap=0.0, - backend="auto", - lse=lse, - ) +def test_auto_backend_writes_lse_on_a_2d_routed_shape(): + """A shape the runner used to refuse: backend='auto' routes it to tiled + 2D, which now writes LSE itself (the split-KV runner is patched to fail).""" + problem, q, kc, vc, cu, su, bt = _case2d([256], [256]) + assert problem.select_path() == "2d" + lse = _nan_lse(problem) + out = _run2d(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0, backend="auto") + off = _run2d_lse_off(problem, q, kc, vc, cu, su, bt, softcap=0.0, backend="auto") + o_ref, lse_ref = _paged_reference( + q, kc, vc, [256], [256], bt, 1.0 / math.sqrt(128), 0.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, "bf16") + + +def _forbid_tiled_2d(*_args, **_kwargs): + raise AssertionError("an LSE request the tiled 2D gate declines reached it") + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize( + "q_lens,kv_lens", + [([2048], [2048]), ([37, 300, 1], [37, 812, 900])], +) +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +def test_auto_hands_lse_to_split_kv_when_tiled_2d_declines(dtype, q_lens, kv_lens): + """gfx942 D256 prefill at block size 32 routes to 2D, but the tiled 2D + gate declines it with LSE on (LDS budget). backend='auto' runs split-KV, + which writes LSE (the tiled 2D launcher is patched to fail).""" + if _arch() != "gfx942": + pytest.skip("the tiled 2D LDS gate declines this shape on gfx942 only") + problem, q, kc, vc, cu, su, bt = _case(q_lens, kv_lens, d=256, dtype=dtype, bs=32) + assert problem.select_path() == "2d" + ok, why = au.supports_native_unified_attention_tiled(problem, "gfx942") + assert not ok and "LDS budget" in why + lse = _nan_lse(problem) + with mock.patch.object(au, "_get_2d_launcher", _forbid_tiled_2d): + out = _run(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0, backend="auto") + off = _run( + dataclasses.replace(problem, emit_lse=False), + q, + kc, + vc, + cu, + su, + bt, + lse=None, + softcap=0.0, + ) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(256), 0.0 + ) + assert torch.equal(out, off), "LSE-on changed the split-KV O" + _assert_matches_reference(out, lse, o_ref, lse_ref, _tolerance(dtype)) @requires_cdna3_or_4 @pytest.mark.gpu def test_cached_lse_off_2d_graph_never_serves_an_lse_request(): """Regression: the gfx942 2D graph fast path replays a cache hit before any - support gate, keyed on a signature without emit_lse. An LSE-off launch on - these exact buffers must not be replayed for an LSE-on request. + support gate, keyed on _cheap_2d_sig (no emit_lse) + buffer ids. An LSE-off + graph captured on these exact buffers must not be replayed for an LSE-on + request: the replay would write O and leave lse unwritten. - On gfx942 the test first proves the hazard is live: the LSE-off launch must - capture a graph whose key the LSE-on request would hit. gfx950 has no 2D - graph path, so there it checks only the plain refusal.""" - from kernels.common import attention_unified as au - - problem, q, kc, vc, cu, su, bt = _case([256], [256]) + On gfx942 the test first proves the hazard is live (the LSE-off launch + captures a graph whose key the LSE-on request would hit). Then the LSE-on + call must write O and LSE, match the reference, and add no graph. gfx950 has + no 2D graph path, so there only the served-and-correct half runs.""" + problem, q, kc, vc, cu, su, bt = _case2d([256], [256]) off_problem = dataclasses.replace(problem, emit_lse=False) out = torch.empty_like(q) args = dict( @@ -380,6 +642,7 @@ def test_cached_lse_off_2d_graph_never_serves_an_lse_request(): run_unified_attention_torch(problem=off_problem, **args) # warm the 2D graph captured = set(au._2D_GRAPHS) - before run_unified_attention_torch(problem=off_problem, **args) # cache hit: replay + torch.cuda.synchronize() if _arch() == "gfx942": assert len(captured) == 1, "LSE-off launch did not capture a 2D graph" (key,) = captured @@ -388,18 +651,20 @@ def test_cached_lse_off_2d_graph_never_serves_an_lse_request(): assert set(au._2D_GRAPHS) - before == captured, "second launch was not a replay" else: assert not captured, "2D graph replay is expected only on gfx942" - torch.cuda.synchronize() - # A replay writes ``out``, never ``lse``: poison ``out`` in place (same - # tensor, so the captured graph key still matches) to see any replay. + off_out = out.clone() + graphs = set(au._2D_GRAPHS) + # A replay writes ``out``, never ``lse``: poison both in place (same + # tensors, so the captured graph key still matches). out.fill_(float("nan")) - lse = torch.full( - (problem.total_q, problem.num_query_heads, 1), float("nan"), device="cuda" - ) - with pytest.raises(NotImplementedError, match="LSE"): + lse = _nan_lse(problem) + with mock.patch.object(au, "_run_3d_tiled", _forbid_split_kv): run_unified_attention_torch(problem=problem, lse=lse, **args) torch.cuda.synchronize() - assert torch.isnan(out).all(), "a 2D kernel ran for a refused LSE request" - assert torch.isnan(lse).all(), "a kernel wrote LSE for a refused request" + assert set(au._2D_GRAPHS) == graphs, "an LSE request captured a 2D graph" + o_ref, lse_ref = _paged_reference( + q, kc, vc, [256], [256], bt, 1.0 / math.sqrt(128), 0.0 + ) + _assert_2d_result(out, off_out, lse, o_ref, lse_ref, "bf16") def _aliases_last_lse_element(lse, like): @@ -414,6 +679,7 @@ def _aliases_last_lse_element(lse, like): @requires_cdna3_or_4 @pytest.mark.gpu +@pytest.mark.parametrize("backend", ["3d", "tiled"]) @pytest.mark.parametrize( "case", [ @@ -428,7 +694,7 @@ def _aliases_last_lse_element(lse, like): "additive_bias", ], ) -def test_malformed_lse_buffer_is_refused_before_launch(case): +def test_malformed_lse_buffer_is_refused_before_launch(case, backend): problem, q, kc, vc, cu, su, bt = _case([4, 2], [3000, 1029]) tq, hq = problem.total_q, problem.num_query_heads lse = torch.full((tq, hq, 1), float("nan"), device="cuda") @@ -473,9 +739,466 @@ def test_malformed_lse_buffer_is_refused_before_launch(case): softmax_scale=1.0 / math.sqrt(q.shape[-1]), block_table=bt, softcap=0.0, - backend="3d", + backend=backend, lse=lse, **bias, ) torch.cuda.synchronize() assert torch.isnan(out).all(), "a kernel ran for a malformed lse buffer" + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("q_lens,kv_lens", _PREFILL_SHAPES) +@pytest.mark.parametrize("d", [64, 128, 256]) +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +@pytest.mark.parametrize("backend", ["auto", "tiled"]) +def test_tiled_2d_output_and_lse(backend, dtype, d, q_lens, kv_lens): + """Arch-native routing. On gfx942 bf16 D256 this is the shape the 4-warp + lean body serves with LSE off (with LSE on it takes the default builder).""" + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens, d=d, dtype=dtype) + if backend == "auto": + assert problem.select_path() == "2d", "auto would not route this to 2D" + lse = _nan_lse(problem) + out = _run2d(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0, backend=backend) + off = _run2d_lse_off(problem, q, kc, vc, cu, su, bt, softcap=0.0, backend=backend) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(d), 0.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("q_lens,kv_lens", _PREFILL_SHAPES) +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +def test_tiled_2d_lse_one_query_head_per_kv_head(dtype, q_lens, kv_lens): + """Hq == Hkv (8:8): NQK = 1, so every CTA row is its own token. backend is + 'tiled' because auto sends the ragged 8:8 shape to split-KV.""" + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens, hq=8, hkv=8, dtype=dtype) + lse = _nan_lse(problem) + out = _run2d(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0, backend="tiled") + off = _run2d_lse_off(problem, q, kc, vc, cu, su, bt, softcap=0.0, backend="tiled") + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(128), 0.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +@pytest.mark.parametrize("window", [64, 4096]) +@pytest.mark.parametrize("backend", ["auto", "tiled"]) +def test_tiled_2d_lse_with_sliding_window(backend, window, dtype): + """D128 sliding window: on gfx942 the LSE-off route is the 4-warp GQA body, + so this is where the LSE-on fallback to the default builder is exercised. Seq 0 + (256 queries over 5000 keys) is cut by both windows; seq 1 (64/64) only by 64.""" + q_lens, kv_lens = [256, 64], [5000, 64] + problem, q, kc, vc, cu, su, bt = _case2d( + q_lens, kv_lens, dtype=dtype, sliding_window=window + ) + lse = _nan_lse(problem) + out = _run2d(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=0.0, backend=backend) + off = _run2d_lse_off(problem, q, kc, vc, cu, su, bt, softcap=0.0, backend=backend) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(128), 0.0, window + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("backend", ["auto", "tiled"]) +def test_tiled_2d_lse_with_softcap(backend): + """Softcap is applied before the max, in log2 units: LSE must include it.""" + q_lens, kv_lens = [37, 300, 1], [37, 812, 900] + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens, softcap=30.0) + lse = _nan_lse(problem) + out = _run2d(problem, q, kc, vc, cu, su, bt, lse=lse, softcap=30.0, backend=backend) + off = _run2d_lse_off(problem, q, kc, vc, cu, su, bt, softcap=30.0, backend=backend) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(128), 30.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, "bf16") + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("q_lens,kv_lens", [([37, 300, 1], [37, 812, 900]), ([5], [3])]) +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +@pytest.mark.parametrize("d", [64, 128]) +@pytest.mark.parametrize("branch", _BRANCHES) +def test_pinned_branch_writes_lse(branch, d, dtype, q_lens, kv_lens): + """Every epilogue branch of the arch's tiled 2D builder, forced by an + explicit tuning spec. Branch B (mfma32) is reachable only this way.""" + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens, d=d, dtype=dtype) + off_problem = dataclasses.replace(problem, emit_lse=False) + lse = _nan_lse(problem) + out = _run2d( + problem, + q, + kc, + vc, + cu, + su, + bt, + lse=lse, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(problem, branch), + ) + off = _run2d( + off_problem, + q, + kc, + vc, + cu, + su, + bt, + lse=None, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(off_problem, branch), + ) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(d), 0.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("branch", ["transposed32", "mfma32", "striped16"]) +def test_dispatcher_2d_lse_end_to_end(branch): + """Request -> pinned 2D tuning spec -> bind_torch(tensors with 'lse') -> + launch. (The generic unified_2d candidate has no Torch binding; executable + dispatch goes through the tuning candidates, as for split-KV.)""" + from dispatch.attention import AttentionRequest, dispatch_attention + + arch = _arch() + q_lens, kv_lens = [256, 256], [256, 700] # total_q = batch * seqlen_q + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens) + + def request(emit_lse): + return AttentionRequest( + batch=2, + nhead_q=32, + nhead_k=8, + seqlen_q=256, + seqlen_k=700, + hdim_q=128, + hdim_v=128, + arch=arch, + mask_type=2, # bottom-right causal, as in _paged_reference + dtype="bf16", + emit_lse=emit_lse, + algorithm="unified_tuning", + spec_id=f"{arch}_u2d_{_BRANCH_VARIANT[arch][branch]}", + ) + + on = dispatch_attention(request(True)) + off = dispatch_attention(request(False)) + assert on.spec.path == "2d" + assert _epilogue_branch(on.spec.kernel_spec) == branch + assert on.spec.kernel_spec.emit_lse is True + assert off.spec.kernel_spec.emit_lse is False + tensors = dict(q=q, k=kc, v=vc, cu_seqlens_q=cu, seqused_k=su, block_table=bt) + out, lse, off_out = torch.empty_like(q), _nan_lse(problem), torch.empty_like(q) + on.bind_torch({**tensors, "out": out, "lse": lse}).launch() + off.bind_torch({**tensors, "out": off_out}).launch() + torch.cuda.synchronize() + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(128), 0.0 + ) + _assert_2d_result(out, off_out, lse, o_ref, lse_ref, "bf16") + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("dtype,score", [("fp16", -20.0), ("bf16", -200.0)]) +@pytest.mark.parametrize("branch", _BRANCHES) +def test_2d_lse_masked_leading_tile_very_negative_scores(branch, dtype, score): + """Sliding window 64 over 256 tokens: some rows of each CTA see a fully + masked first KV tile. Every visible score is `score` (natural log). Every + row sees keys, so every LSE must be finite (about score + ln 64) and O must + match the reference, on every epilogue branch.""" + problem, q, kc, vc, cu, su, bt = _case2d( + [256], [256], dtype=dtype, sliding_window=64 + ) + a = 1.0 if dtype == "fp16" else 4.0 + q.fill_(a) + kc.fill_(score / (a * math.sqrt(128))) + lse = _nan_lse(problem) + out = _run2d( + problem, + q, + kc, + vc, + cu, + su, + bt, + lse=lse, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(problem, branch), + ) + off_problem = dataclasses.replace(problem, emit_lse=False) + off = _run2d( + off_problem, + q, + kc, + vc, + cu, + su, + bt, + lse=None, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(off_problem, branch), + ) + o_ref, lse_ref = _paged_reference( + q, kc, vc, [256], [256], bt, 1.0 / math.sqrt(128), 0.0, sliding_window=64 + ) + assert torch.isfinite(lse_ref).all(), "the case must give every row a key" + assert torch.isfinite(lse).all(), "a row with visible keys got LSE -inf/NaN" + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +_NAN_ROW, _NAN_HEAD = 100, 5 + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +@pytest.mark.parametrize("branch", _BRANCHES) +def test_2d_lse_nan_query_row_stays_nan(branch, dtype): + """A NaN in one query row gives NaN O and NaN LSE for that (row, head) + only. Its softmax sum l is NaN; !(l == 0) keeps row_has_mass true, so + lse_store propagates NaN instead of writing the keyless -inf.""" + problem, q, kc, vc, cu, su, bt = _case2d([256], [256], dtype=dtype) + q[_NAN_ROW, _NAN_HEAD] = float("nan") + lse = torch.full( + (problem.total_q, problem.num_query_heads, 1), float("inf"), device="cuda" + ) + out = _run2d( + problem, + q, + kc, + vc, + cu, + su, + bt, + lse=lse, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(problem, branch), + ) + off_problem = dataclasses.replace(problem, emit_lse=False) + off = _run2d( + off_problem, + q, + kc, + vc, + cu, + su, + bt, + lse=None, + softcap=0.0, + backend="tiled", + tuning_spec=_branch_spec(off_problem, branch), + ) + o_ref, lse_ref = _paged_reference( + q, kc, vc, [256], [256], bt, 1.0 / math.sqrt(q.shape[-1]), 0.0 + ) + assert not torch.isposinf(lse).any(), "an LSE row was left unwritten" + torch.testing.assert_close(out, off, rtol=0, atol=0, equal_nan=True) + assert torch.isnan(lse[_NAN_ROW, _NAN_HEAD]).all(), "NaN row got a non-NaN LSE" + assert torch.isnan(out[_NAN_ROW, _NAN_HEAD]).all(), "NaN row was hidden in O" + rest = torch.ones( + problem.total_q, problem.num_query_heads, dtype=torch.bool, device="cuda" + ) + rest[_NAN_ROW, _NAN_HEAD] = False + assert torch.isfinite(lse.squeeze(-1)[rest]).all(), "NaN leaked into another LSE" + assert torch.isfinite(out[rest]).all(), "NaN leaked into another (row, head)" + assert (out[rest].float() - o_ref[rest]).abs().max().item() < _tolerance(dtype) + torch.testing.assert_close( + lse.squeeze(-1)[rest], lse_ref.squeeze(-1)[rest], atol=_LSE_ATOL, rtol=_LSE_RTOL + ) + + +# (variant id, geometry suffix, knobs, epilogue branch). Per arch, so every +# collected case runs: gfx950 has no q-major grid. +_ROW_MAP_VARIANTS = { + "gfx942": [ + ( + "qmajor_transposed32", + "transposed_x8_nw1_mw32_t2xb_llvm", + {"use_q_major_grid": True}, + "transposed32", + ), + ( + "qmajor_transposed32_s1", + "transposed_x8_nw1_mw32_t2xb_llvm", + {"use_q_major_grid": True, "use_transposed_scalar_state": True}, + "transposed32", + ), + ( + "qmajor_striped16", + "narrow_nw1_mw16_t1xb_llvm", + {"use_q_major_grid": True}, + "striped16", + ), + ("striped16_mw32", "narrow_nw1_mw32_t1xb_llvm", {}, "striped16"), + ("striped16_mw32_nw2", "narrow_nw2_mw32_t1xb_llvm", {}, "striped16"), + ], + "gfx950": [ + ("striped16_mw32", "narrow_nw1_mw32_t1xb_llvm", {}, "striped16"), + ("striped16_mw32_nw2", "narrow_nw2_mw32_t1xb_llvm", {}, "striped16"), + ], +} + + +def _variant_spec(problem, geometry, knobs, branch): + from dispatch.attention.unified_rules import canonicalize_tuning_spec + + arch = _arch() + if arch == "gfx942": + from dispatch.attention.gfx942_unified import ( + GFX942_TUNING_VARIANTS as variants, + ) + else: + from dispatch.attention.gfx950_unified import ( + GFX950_TUNING_VARIANTS as variants, + ) + spec_id = f"{arch}_u2d_{geometry}" + variant = next(v for v in variants if v.spec_id == spec_id) + spec, why = canonicalize_tuning_spec(problem, variant, knobs) + assert spec is not None, f"{spec_id} {knobs}: {why}" + ks = spec.kernel_spec + assert _epilogue_branch(ks) == branch, f"{spec_id} selects {_epilogue_branch(ks)}" + for name, value in knobs.items(): + assert getattr(ks, name) == value, f"{spec_id}: {name} did not stick" + assert ks.emit_lse is problem.emit_lse + return spec + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize("q_lens,kv_lens", [([37, 300, 1], [37, 812, 900]), ([5], [3])]) +@pytest.mark.parametrize("dtype", ["fp16", "bf16"]) +@pytest.mark.parametrize( + "vid,geometry,knobs,branch", + _ROW_MAP_VARIANTS.get(_arch(), []), + ids=[v[0] for v in _ROW_MAP_VARIANTS.get(_arch(), [])], +) +def test_row_mapping_variant_writes_lse( + vid, geometry, knobs, branch, dtype, q_lens, kv_lens +): + """Swapped grid, two atoms per warp and a second warp change which lane + and slot owns each (token, head). A wrong mapping writes the right values + to the wrong index, which only a reference comparison sees. The ragged + shape spans three sequences; [5]/[3] has keyless rows 0-1.""" + problem, q, kc, vc, cu, su, bt = _case2d(q_lens, kv_lens, dtype=dtype) + off_problem = dataclasses.replace(problem, emit_lse=False) + lse = _nan_lse(problem) + out = _run2d( + problem, + q, + kc, + vc, + cu, + su, + bt, + lse=lse, + softcap=0.0, + backend="tiled", + tuning_spec=_variant_spec(problem, geometry, knobs, branch), + ) + off = _run2d( + off_problem, + q, + kc, + vc, + cu, + su, + bt, + lse=None, + softcap=0.0, + backend="tiled", + tuning_spec=_variant_spec(off_problem, geometry, knobs, branch), + ) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, 1.0 / math.sqrt(q.shape[-1]), 0.0 + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, dtype) + + +@requires_cdna3_or_4 +@pytest.mark.gpu +@pytest.mark.parametrize( + "q_lens,kv_lens,branch", + [ + ([37, 300, 1], [37, 812, 900], "transposed32"), + ([5], [3], "striped16"), + ], +) +def test_tiled_2d_lse_with_additive_bias(q_lens, kv_lens, branch): + """Additive bias is added to the log2 score before the running max, so LSE + must include it. gfx950 default routing: the ragged shape takes branch A + (16 slots: bias turns scalar state off), 5/3 takes C and keeps two keyless + rows. gfx942 has no bias kernel: the tiled 2D gate must refuse it.""" + problem, q, kc, vc, cu, su, bt = _case(q_lens, kv_lens) + problem = dataclasses.replace(problem, use_additive_bias=True) + arch = _arch() + ok, why = au.supports_native_unified_attention_tiled(problem, arch) + if arch == "gfx942": + assert not ok and "only implemented on gfx950" in why, why + return + assert ok, f"the tiled 2D gate refuses LSE + bias: {why}" + assert _epilogue_branch(au._tiled_spec_from_problem(problem, arch)) == branch + g = torch.Generator(device="cuda").manual_seed(1) + bias = torch.randn( + len(q_lens), + 32, + problem.max_seqlen_q, + problem.max_seqlen_k, + device="cuda", + generator=g, + ) + bias_args = dict( + additive_bias=bias, + additive_bias_batch_stride=bias.stride(0), + additive_bias_head_stride=bias.stride(1), + additive_bias_sq_stride=bias.stride(2), + ) + scale = 1.0 / math.sqrt(128) + + def launch(prob, lse): + out = torch.empty_like(q) + with mock.patch.object(au, "_run_3d_tiled", _forbid_split_kv): + run_unified_attention_torch( + problem=prob, + q=q, + k=kc, + v=vc, + out=out, + cu_seqlens_q=cu, + seqused_k=su, + softmax_scale=scale, + block_table=bt, + softcap=0.0, + backend="tiled", + lse=lse, + **bias_args, + ) + torch.cuda.synchronize() + return out + + lse = _nan_lse(problem) + out = launch(problem, lse) + off = launch(dataclasses.replace(problem, emit_lse=False), None) + o_ref, lse_ref = _paged_reference( + q, kc, vc, q_lens, kv_lens, bt, scale, 0.0, additive_bias=bias + ) + _assert_2d_result(out, off, lse, o_ref, lse_ref, "bf16") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lse_store_ir.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lse_store_ir.py index 9c00f8173172..f9bd50fefcee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lse_store_ir.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lse_store_ir.py @@ -6,6 +6,8 @@ from __future__ import annotations import re +import shutil +import subprocess import pytest from ._lse_store_support import factory_for @@ -13,6 +15,7 @@ import _attention_fwd_ext_harness as H from kernels.common import _lse_store as L from rocke.core.ir import F32, I32, IRBuilder, PtrType +from rocke.core.lower_hip import lower_kernel_to_hip from rocke.core.lower_llvm import _lower_kernel_to_llvm_python as lower from rocke.helpers.attention_fwd_ext import RowEpilogue @@ -130,3 +133,51 @@ def test_batch_rebase_is_emitted_once_with_a_64_bit_offset(): def test_no_exported_helper_looks_like_a_builder(): assert not [n for n in L.__all__ if n.startswith("build_")] + + +# Host harness for the bare HIP kernel: the body is plain C++ once the HIP +# function qualifiers are defined away, so it runs on the CPU. +_HIP_HOST_PRELUDE = """#include +#include +using std::isnan; +#define __global__ +#define __launch_bounds__(n) +""" +_HIP_HOST_MAIN = """ +int main() { + const float xs[] = {0.0f, -0.0f, 1.0f, -3.5f, NAN, INFINITY}; + for (float x : xs) { + float r = -1.0f; + lse_row_has_mass_probe(&r, x); + std::printf("%g\\n", r); + } +} +""" + + +def test_row_has_mass_is_not_l_equals_zero_in_the_hip_lowering(tmp_path): + """lse_row_has_mass must be !(l == 0) with NaN counted as mass, in the HIP + lowerer too (the GPU tests only exercise the LLVM path). A NaN row must keep + NaN, so the predicate is true for NaN; only an exact zero (either sign) is + keyless.""" + compiler = shutil.which("c++") or shutil.which("g++") or shutil.which("clang++") + if compiler is None: + pytest.skip("no host C++ compiler") + b = IRBuilder("lse_row_has_mass_probe") + out = b.param("out", PtrType(F32, "global")) + l = b.param("l", F32) + has_mass = L.lse_row_has_mass(b, l) + one, zero = b.const_f32(1.0), b.const_f32(0.0) + b.global_store(out, b.const_i32(0), b.select(has_mass, one, zero), align=4) + b.ret() + src = tmp_path / "probe.cpp" + src.write_text( + _HIP_HOST_PRELUDE + + lower_kernel_to_hip(b.kernel, include_prologue=False) + + _HIP_HOST_MAIN + ) + exe = tmp_path / "probe" + subprocess.run([compiler, "-O2", "-o", str(exe), str(src)], check=True) + got = subprocess.run([str(exe)], check=True, capture_output=True, text=True) + # 0 -0 1 -3.5 NaN inf + assert got.stdout.split() == ["0", "0", "1", "1", "1", "1"] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index b38762287e96..fd67459212a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -645,8 +645,8 @@ if(ROCKE_INSTALL_TESTS) # Library attention CPU tests (no GPU required): gfx950 LDS budget model, # dense block-order decodes, spec-wrapper parity gate, the 3D decode split rule # and K/V cache policy, the scalar fallback's cache key and block-table - # stride kernarg, and the unified attention running-max carry and keyless - # guards on the emitted IR. + # stride kernarg, the unified attention running-max carry and keyless + # guards on the emitted IR, and the tiled 2D LSE store on the emitted IR. # File-list style (same pattern as conv_pytest) so PYTHONPATH resolves rocke/ # correctly from the installed tree. set(_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS @@ -656,7 +656,8 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_attention_spec_predicate_parity.py" "tests/library/tests/test_attention_tiled_3d_decode_split.py" "tests/library/tests/test_attention_tiled_3d_kv_cache_policy.py" - "tests/library/tests/test_attention_unified_masked_tile_ir.py") + "tests/library/tests/test_attention_unified_masked_tile_ir.py" + "tests/library/tests/test_attention_unified_2d_lse_ir.py") set(_ROCKE_LIBRARY_ATTENTION_CPU_TESTS "") foreach(_t ${_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS}) string(APPEND _ROCKE_LIBRARY_ATTENTION_CPU_TESTS diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py new file mode 100644 index 000000000000..b15211806f7f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py @@ -0,0 +1,1880 @@ +"""Parity configs for the tiled attention families of prove_parity_binding.py. + +Each table mirrors ``library/tests/parity/_emit.{py,c}`` index for +index (the C emitter is what the binding compares against). The 3D tables set +``emit_lse`` on the configs whose reduce the emitter builds with LSE +(``_REDUCE_LSE``). The platform layer keeps its own copy: it must not import +``library`` (one-way dependency, library -> platform). prove_parity_binding +fails a family unless its C emitter reports "unknown config index N" for the +first index N past the table, and a row that drifts from the emitter shows up +as a DIFF. +""" + +GFX942_ATTENTION_TILED_2D = [ + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 2048, + "has_softcap": False, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + }, + { + "head_size": 64, + "block_size": 128, + "num_query_heads": 16, + "num_kv_heads": 4, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 4096, + "has_softcap": False, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 2048, + "has_softcap": True, + }, + { + "head_size": 256, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 8, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32x8": True, + "use_transposed_qk_32x32": True, + "use_k_single_buffer": False, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 8, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32x8": True, + "use_transposed_qk_32x32": True, + "use_k_single_buffer": False, + "use_agpr_alloc_zero": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 256, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 64, + "use_i64_kv_addr": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 256, + "num_warps": 4, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 64, + "use_i64_kv_addr": True, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 256, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 64, + "use_i64_kv_addr": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 256, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 64, + "use_i64_kv_addr": True, + "use_global_load_lds_k": True, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 2, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 32, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 2, + "waves_per_eu": 2, + "tile_size": 64, + "block_m_per_warp": 32, + "use_mfma_32x32x8": True, + "use_transposed_qk_32x32": True, + "use_k_single_buffer": True, + "kv_cache_policy": "all", + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 8, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32x8": True, + "use_transposed_qk_32x32": True, + "use_k_single_buffer": False, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 12, + "num_kv_heads": 2, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32x8": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_k_single_buffer": False, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 12, + "num_kv_heads": 2, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32x8": True, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 12, + "num_kv_heads": 2, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "block_m_per_warp": 32, + "use_q_major_grid": True, + "emit_lse": True, + }, +] + +GFX950_ATTENTION_TILED_2D = [ + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 1, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 1, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 256, + "block_size": 16, + "num_query_heads": 1, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 2, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 2048, + "has_softcap": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 1, + "has_softcap": False, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + }, + { + "head_size": 256, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 64, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 7, + "num_kv_heads": 7, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 40, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 128, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_qq_bias": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_alibi": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 512, + "has_softcap": True, + "use_alibi": True, + "use_qq_bias": True, + }, + { + "head_size": 64, + "block_size": 64, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 8, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "tile_size": 128, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "tile_size": 128, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "waves_per_eu": 2, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 257, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": "fp8e4m3", + "use_fp8_mfma_pv": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_i64_kv_addr": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_register_pv": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": "fp8e4m3", + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": "fp8e4m3", + "use_fp8_mfma_qk": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_grouped_kv2_softmax": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_grouped_kv2_softmax": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_grouped_kv2_softmax": True, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_early_v_schedule": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 16, + "tile_size": 64, + "use_fast_paged_kv_desc": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "num_warps": 4, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_fast_paged_kv_desc": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_transposed_half_local_pv": True, + "use_early_v_schedule": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_grouped_kv2_softmax": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_early_v_schedule": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 1, + "block_m_per_warp": 32, + "tile_size": 64, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_sched_barrier": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 2, + "block_m_per_warp": 32, + "tile_size": 64, + "use_k_single_buffer": True, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "block_m_per_warp": 32, + "tile_size": 128, + "use_k_single_buffer": True, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_invariant_hoist": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_softmax_mfma_interleave": True, + "softmax_interleave_mode": 1, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "num_warps": 4, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "num_warps": 4, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_fast_paged_kv_desc": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "kv_storage_dtype": "fp8e4m3", + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 64, + "block_size": 64, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "kv_storage_dtype": "fp8e4m3", + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "block_m_per_warp": 32, + "tile_size": 64, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 2, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 32, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "tile_size": 64, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_half_local_pv": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_mask_limit": True, + "use_softmax_mfma_interleave": True, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 128, + "num_kv_heads": 8, + "dtype": "fp16", + "use_sinks": True, + "sliding_window": 128, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "tile_size": 64, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_half_local_pv": True, + "use_softmax_mfma_interleave": True, + }, + { + "head_size": 128, + "block_size": 64, + "num_query_heads": 128, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "tile_size": 128, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_half_local_pv": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_mask_limit": True, + "use_k_single_buffer": True, + "use_softmax_mfma_interleave": True, + }, + { + "head_size": 64, + "block_size": 64, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 128, + "has_softcap": False, + "num_seqs": 1, + "num_warps": 4, + "waves_per_eu": 2, + "tile_size": 128, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_half_local_pv": True, + "use_early_v_schedule": True, + }, + { + "head_size": 256, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 2, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "use_qblock_reverse": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "num_warps": 4, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_mask_limit": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_half_local_pv": True, + "use_fast_paged_kv_desc": True, + "block_m_per_warp": 32, + "tile_size": 64, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "fp16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_mfma_32x32": True, + "block_m_per_warp": 32, + "tile_size": 64, + "emit_lse": True, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 1, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "emit_lse": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 32, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "use_alibi": True, + "emit_lse": True, + }, +] + +GFX942_ATTENTION_TILED_3D = [ + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 8, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 16, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 256, + "block_size": 64, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 4096, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + "kv_storage_dtype": "fp8e4m3", + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 2048, + "has_softcap": False, + "use_alibi": True, + "use_qq_bias": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 1, + "dtype": "bf16", + "num_segments": 32, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": "fp8e4m3", + "use_decode_grid": True, + }, + { + "head_size": 256, + "block_size": 64, + "num_query_heads": 16, + "num_kv_heads": 2, + "dtype": "bf16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 32, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_wide_kv_load": True, + "use_decode_grid": True, + "kv_cache_policy": "stream", + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "kv_cache_policy": "nt", + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_wide_kv_load": True, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + "kv_storage_dtype": None, + "use_wide_kv_load": True, + "use_kv_lds_swizzle": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_wide_kv_load": True, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + "use_wave_page_id_load": True, + }, + { + "head_size": 64, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 64, + "has_softcap": False, + "kv_storage_dtype": None, + "use_wide_kv_load": True, + "use_wave_page_id_load": True, + }, + { + "head_size": 256, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 64, + "dtype": "bf16", + "num_segments": 32, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + }, + { + "head_size": 256, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 64, + "dtype": "bf16", + "num_segments": 32, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + "use_wave_page_id_load": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 16, + "use_wide_kv_load": True, + "use_wave_page_id_load": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 8, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 16, + "use_wave_page_id_load": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 40, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_segments": 16, + "num_seqs": 1, + "use_decode_grid": True, + "kv_cache_policy": "stream", + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "emit_lse": True, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 8, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "emit_lse": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "waves_per_eu": 2, + }, +] + +GFX950_ATTENTION_TILED_3D = [ + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 8, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 16, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 256, + "block_size": 64, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 4096, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + "kv_storage_dtype": "fp8e4m3", + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 2048, + "has_softcap": False, + "use_alibi": True, + "use_qq_bias": False, + "kv_storage_dtype": None, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 16, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_i64_kv_addr": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + }, + { + "head_size": 128, + "block_size": 32, + "num_query_heads": 16, + "num_kv_heads": 1, + "dtype": "bf16", + "num_segments": 32, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": "fp8e4m3", + "use_decode_grid": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 128, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 32, + "use_decode_grid": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 64, + "use_i64_kv_addr": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "kv_cache_policy": "stream", + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 32, + "use_i64_kv_addr": True, + "use_decode_grid": True, + "kv_cache_policy": "stream", + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 64, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 64, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + }, + { + "head_size": 256, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 2, + "dtype": "fp16", + "num_segments": 16, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": True, + "kv_storage_dtype": None, + "use_i64_kv_addr": True, + "use_kv_lds_swizzle": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 64, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_wave_page_id_load": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 4, + "dtype": "fp16", + "num_segments": 64, + "use_sinks": True, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "tile_size_override": 64, + "use_i64_kv_addr": True, + "use_decode_grid": True, + "kv_cache_policy": "stream", + "use_kv_lds_swizzle": True, + "use_wave_page_id_load": True, + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 40, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 96, + "num_kv_heads": 96, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_segments": 16, + "num_seqs": 1, + "use_decode_grid": True, + "kv_cache_policy": "stream", + }, + { + "head_size": 128, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "num_segments": 16, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "emit_lse": True, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 8, + "num_kv_heads": 8, + "dtype": "fp16", + "num_segments": 8, + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "kv_storage_dtype": None, + "emit_lse": True, + }, +] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/prove_parity_binding.py b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/prove_parity_binding.py index da8ccceef172..75888b8de218 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/prove_parity_binding.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/prove_parity_binding.py @@ -42,6 +42,20 @@ sys.path.insert(0, BUILD) import rocke_engine # noqa: E402 +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import attention_tiled_parity_configs as atp # noqa: E402 + +# Families whose table must cover every emitter config: the emitter must report +# "unknown config index N" for the first index N past the table. +_COMPLETE_TABLES = frozenset( + { + "gfx942_attention_tiled_2d", + "gfx950_attention_tiled_2d", + "gfx942_attention_tiled_3d", + "gfx950_attention_tiled_3d", + } +) + def _assert_build_id_match(): """Cross-check the loaded .so's build-id against the emit dir's BUILD_ID. @@ -1408,188 +1422,11 @@ def cfgs_fused_moe_e2e(): def cfgs_gfx942_attention_tiled_2d(): - rows = [ - (64, 32, 32, 32, "bf16", True, 2048, False), - (128, 32, 32, 32, "fp16", True, 0, True), - (64, 128, 16, 4, "bf16", False, 4096, False), - (128, 64, 32, 8, "fp16", True, 2048, True), - (256, 32, 32, 32, "bf16", False, 0, False), - ] - return [ - dict( - head_size=h, - block_size=b, - num_query_heads=q, - num_kv_heads=k, - dtype=d, - use_sinks=s, - sliding_window=sw, - has_softcap=sc, - ) - for (h, b, q, k, d, s, sw, sc) in rows - ] + return [dict(c) for c in atp.GFX942_ATTENTION_TILED_2D] def cfgs_gfx950_attention_tiled_2d(): - """Full 38-config table mirroring gfx950_attention_tiled_2d_emit.c make_spec.""" - - def base(h, b, q, k, dt, sink, sw, sc, **extra): - d = dict( - head_size=h, - block_size=b, - num_query_heads=q, - num_kv_heads=k, - dtype=dt, - use_sinks=sink, - sliding_window=sw, - has_softcap=sc, - ) - d.update(extra) - return d - - return [ - base(64, 16, 1, 1, "bf16", False, 0, False), # 0 - base(128, 16, 1, 1, "bf16", False, 0, False), # 1 - base(256, 16, 1, 1, "bf16", False, 0, False), # 2 - base(64, 16, 16, 1, "bf16", False, 0, False), # 3 - base(64, 16, 2, 1, "bf16", False, 0, False), # 4 - base(64, 32, 32, 32, "bf16", False, 0, False), # 5 - base(64, 32, 32, 32, "fp16", True, 2048, True), # 6 - base(64, 32, 32, 32, "bf16", True, 1, False), # 7 - base(128, 32, 32, 32, "fp16", True, 0, True), # 8 - base(256, 32, 32, 32, "bf16", False, 0, False), # 9 - base(64, 64, 32, 32, "bf16", False, 0, False), # 10 - base(64, 32, 7, 7, "bf16", False, 0, False), # 11 - base(64, 32, 64, 8, "bf16", False, 0, False), # 12 - base(64, 32, 40, 8, "bf16", False, 0, False), # 13 - base(64, 32, 128, 1, "bf16", False, 0, False), # 14 - base(64, 32, 32, 32, "bf16", False, 0, False, use_qq_bias=True), # 15 - base(64, 32, 32, 32, "bf16", False, 0, False, use_alibi=True), # 16 - base( - 64, 32, 32, 32, "fp16", True, 512, True, use_alibi=True, use_qq_bias=True - ), # 17 - base(64, 64, 64, 8, "bf16", False, 0, False, num_warps=8), # 18 - base(64, 32, 32, 32, "bf16", False, 0, False, num_warps=2), # 19 - base(64, 32, 64, 8, "bf16", False, 0, False, num_warps=4), # 20 - base(64, 32, 32, 32, "bf16", False, 0, False, tile_size=64), # 21 - base(64, 32, 32, 32, "bf16", False, 0, False, tile_size=128), # 22 - base( - 128, 32, 32, 32, "fp16", False, 0, False, num_warps=2, tile_size=128 - ), # 23 - base(64, 32, 32, 32, "bf16", False, 0, False, waves_per_eu=2), # 24 - base(64, 32, 32, 32, "bf16", False, 0, False, num_seqs=1), # 25 - base(64, 32, 32, 32, "bf16", False, 0, False, num_seqs=257), # 26 - base( - 64, - 32, - 32, - 32, - "bf16", - False, - 0, - False, - kv_storage_dtype="fp8e4m3", - use_fp8_mfma_pv=True, - ), # 27 - base(64, 32, 32, 32, "bf16", False, 0, False, use_i64_kv_addr=True), # 28 - base(64, 32, 32, 32, "bf16", False, 0, False, use_register_pv=True), # 29 - base(64, 32, 32, 32, "bf16", False, 0, False, kv_storage_dtype="fp8e4m3"), # 30 - base( - 64, - 32, - 32, - 32, - "bf16", - False, - 0, - False, - kv_storage_dtype="fp8e4m3", - use_fp8_mfma_qk=True, - ), # 31 - base( - 128, - 32, - 32, - 32, - "fp16", - False, - 0, - False, - use_mfma_32x32=True, - block_m_per_warp=32, - tile_size=64, - ), # 32 - base( - 64, - 32, - 32, - 32, - "bf16", - False, - 0, - False, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_transposed_scalar_state=True, - use_transposed_invariant_hoist=True, - use_transposed_mask_once=True, - use_grouped_kv2_softmax=True, - block_m_per_warp=32, - tile_size=64, - ), # 33 - base( - 128, - 32, - 32, - 32, - "bf16", - False, - 0, - False, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_transposed_scalar_state=True, - use_transposed_invariant_hoist=True, - use_transposed_mask_once=True, - use_grouped_kv2_softmax=True, - block_m_per_warp=32, - tile_size=64, - ), # 34 - base( - 64, - 32, - 64, - 8, - "bf16", - False, - 0, - False, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_transposed_scalar_state=True, - use_transposed_invariant_hoist=True, - use_transposed_mask_once=True, - use_grouped_kv2_softmax=True, - num_warps=4, - block_m_per_warp=32, - tile_size=64, - ), # 35 - base(64, 32, 32, 32, "bf16", False, 0, False, use_early_v_schedule=True), # 36 - base( - 64, - 32, - 64, - 8, - "bf16", - False, - 0, - False, - num_warps=4, - block_m_per_warp=16, - tile_size=64, - use_fast_paged_kv_desc=True, - ), # 37 - ] + return [dict(c) for c in atp.GFX950_ATTENTION_TILED_2D] def cfgs_gfx950_attention_tiled_2d_fastkv_regp(): @@ -1602,6 +1439,7 @@ def cfgs_gfx950_attention_tiled_2d_fastkv_regp(): {"use_transposed_half_local_pv": True, "use_mfma32_skip_legacy_qreg": True}, {"use_agpr_alloc_zero": True, "use_transposed_half_local_pv": True}, {"use_grouped_kv2_softmax": True}, + {"emit_lse": True}, ] out = [] for dl in deltas: @@ -1956,36 +1794,11 @@ def run_selector_parity(arch, problems): def cfgs_gfx942_attention_tiled_3d(): - rows = [ - (64, 16, 8, 8, "fp16", 8, False, 0, False, None, False, False), - (128, 32, 16, 16, "bf16", 16, False, 0, False, None, False, False), - (256, 64, 32, 8, "fp16", 8, False, 4096, False, None, False, False), - (128, 32, 32, 8, "bf16", 16, True, 0, True, "fp8e4m3", False, False), - (64, 16, 16, 4, "fp16", 8, False, 2048, False, None, True, False), - ] - out = [] - for h, b, q, k, dt, seg, sink, sw, sc, kv, alibi, qqb in rows: - d = dict( - head_size=h, - block_size=b, - num_query_heads=q, - num_kv_heads=k, - dtype=dt, - num_segments=seg, - use_sinks=sink, - sliding_window=sw, - has_softcap=sc, - use_alibi=alibi, - use_qq_bias=qqb, - ) - if kv is not None: - d["kv_storage_dtype"] = kv - out.append(d) - return out + return [dict(c) for c in atp.GFX942_ATTENTION_TILED_3D] def cfgs_gfx950_attention_tiled_3d(): - return cfgs_gfx942_attention_tiled_3d() + return [dict(c) for c in atp.GFX950_ATTENTION_TILED_3D] def cfgs_deep_fused_conv_pool(): @@ -2228,6 +2041,19 @@ def main(): f" cfg{idx}: DIFF cpp={sha(got)} std={sha(std)} " f"(len {len(got)} vs {len(std)})" ) + if fam in _COMPLETE_TABLES: + # The first index past the table must be unknown to the emitter. A + # build (the table is short), a reject or an exec error all fail. + end = len(cfgs) + past, past_err = run_emitter(fam, end) + if past != "" or f"unknown config index {end}" not in (past_err or ""): + n_mismatch += 1 + why = "builds it" if past else f"says {past_err!r}" + details.append( + f" cfg{end}: the emitter {why}, not 'unknown config index " + f"{end}'; attention_tiled_parity_configs.py must cover every " + "emitter config" + ) total_match += n_match total_mismatch += n_mismatch if n_nonempty == 0: diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine_attention.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine_attention.cpp index b1094cd5e156..381aab312aaf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine_attention.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/rocke_engine_attention.cpp @@ -823,11 +823,56 @@ rocke_attention_tiled_2d_spec_t t2d_build(const py::dict& d, Store& st) s.use_early_v_schedule = a_bool(d, "use_early_v_schedule", s.use_early_v_schedule); s.use_fast_paged_kv_desc = a_bool(d, "use_fast_paged_kv_desc", s.use_fast_paged_kv_desc); s.use_qblock_reverse = a_bool(d, "use_qblock_reverse", s.use_qblock_reverse); + s.emit_lse = a_bool(d, "emit_lse", s.emit_lse); + /* Every other field the C struct shares with the Python spec. A field left + * out here silently keeps its C default, so the native builder would build + * a different kernel (e.g. gfx942 x8 transposed specs fail validation + * without use_mfma_32x32x8). */ + s.use_mfma_32x32x8 = a_bool(d, "use_mfma_32x32x8", s.use_mfma_32x32x8); + s.use_q_major_grid = a_bool(d, "use_q_major_grid", s.use_q_major_grid); + s.use_transposed_half_local_pv + = a_bool(d, "use_transposed_half_local_pv", s.use_transposed_half_local_pv); + s.use_transposed_mask_limit + = a_bool(d, "use_transposed_mask_limit", s.use_transposed_mask_limit); + s.use_mfma32_skip_legacy_qreg + = a_bool(d, "use_mfma32_skip_legacy_qreg", s.use_mfma32_skip_legacy_qreg); + s.use_agpr_alloc_zero = a_bool(d, "use_agpr_alloc_zero", s.use_agpr_alloc_zero); + s.use_k_single_buffer = a_bool(d, "use_k_single_buffer", s.use_k_single_buffer); + s.use_conflict_free_v = a_bool(d, "use_conflict_free_v", s.use_conflict_free_v); + s.use_conflict_free_v_ck_vlds + = a_bool(d, "use_conflict_free_v_ck_vlds", s.use_conflict_free_v_ck_vlds); + s.use_conflict_free_v_store + = a_bool(d, "use_conflict_free_v_store", s.use_conflict_free_v_store); + s.use_conflict_free_v_store_split + = a_bool(d, "use_conflict_free_v_store_split", s.use_conflict_free_v_store_split); + s.use_global_load_lds_k = a_bool(d, "use_global_load_lds_k", s.use_global_load_lds_k); + s.use_iglp_opt = a_bool(d, "use_iglp_opt", s.use_iglp_opt); + s.use_k_sliced_ldsseq = a_bool(d, "use_k_sliced_ldsseq", s.use_k_sliced_ldsseq); + s.use_k_sliced_ring = a_bool(d, "use_k_sliced_ring", s.use_k_sliced_ring); + s.use_q_direct_global = a_bool(d, "use_q_direct_global", s.use_q_direct_global); + s.use_qk_pv_sched_group_barrier + = a_bool(d, "use_qk_pv_sched_group_barrier", s.use_qk_pv_sched_group_barrier); + s.k_slice_hd = a_int(d, "k_slice_hd", s.k_slice_hd); + s.ring_depth = a_int(d, "ring_depth", s.ring_depth); + s.use_q_direct_reg = a_bool(d, "use_q_direct_reg", s.use_q_direct_reg); + s.use_q_reread = a_bool(d, "use_q_reread", s.use_q_reread); + s.use_sched_barrier = a_bool(d, "use_sched_barrier", s.use_sched_barrier); + s.use_softmax_mfma_interleave + = a_bool(d, "use_softmax_mfma_interleave", s.use_softmax_mfma_interleave); + s.use_staggered_iter_wait = a_bool(d, "use_staggered_iter_wait", s.use_staggered_iter_wait); + s.use_v_double_buffer = a_bool(d, "use_v_double_buffer", s.use_v_double_buffer); + s.kv_ring_depth = a_int(d, "kv_ring_depth", s.kv_ring_depth); + s.sched_barrier_mask = a_int(d, "sched_barrier_mask", s.sched_barrier_mask); + s.softmax_interleave_groups + = a_int(d, "softmax_interleave_groups", s.softmax_interleave_groups); + s.softmax_interleave_mode = a_int(d, "softmax_interleave_mode", s.softmax_interleave_mode); std::string v; if(a_str(d, "dtype", v)) s.dtype = st.keep(v); if(a_str(d, "kv_storage_dtype", v)) s.kv_storage_dtype = st.keep(v); + if(a_str(d, "kv_cache_policy", v)) + s.kv_cache_policy = st.keep(v); return s; } @@ -984,6 +1029,7 @@ rocke_attention_tiled_2d_spec_t fkv_build(const py::dict& d, Store& st) = a_bool(d, "use_mfma32_skip_legacy_qreg", s.use_mfma32_skip_legacy_qreg); s.use_agpr_alloc_zero = a_bool(d, "use_agpr_alloc_zero", s.use_agpr_alloc_zero); s.use_grouped_kv2_softmax = a_bool(d, "use_grouped_kv2_softmax", s.use_grouped_kv2_softmax); + s.emit_lse = a_bool(d, "emit_lse", s.emit_lse); std::string v; if(a_str(d, "dtype", v)) s.dtype = st.keep(v); @@ -1109,16 +1155,16 @@ bool t3d_strided(const py::dict& d) return layout == "strided"; } -/* gfx942 3d: emitter lowers only the segment via the family convenience. */ -std::string t3d942_lower(const py::dict& d, const std::string& arch) +/* gfx942 3d: the emitter prints the segment, then the reduce derived from it + * (head_size / num_query_heads / num_kv_heads / dtype / num_segments, plus + * emit_lse), concatenated. */ +static std::string t3d942_segment_ll(const py::dict& d, + const rocke_unified_attention_3d_tiled_spec_t& s, + const char* target) { - require_arch(arch); - Store st; - rocke_unified_attention_3d_tiled_spec_t s = t3d_build(d, st); char* ll = nullptr; char err[ROCKE_ERR_MSG_CAP]; err[0] = '\0'; - const char* target = arch.c_str(); if(!t3d_strided(d)) { rocke_status_t status = rocke_build_unified_attention_3d_tiled_gfx942_lower_to_llvm( @@ -1140,6 +1186,32 @@ std::string t3d942_lower(const py::dict& d, const std::string& arch) guard.reset(); return take_ll(status, ll, err, "rocke_engine.gfx942_attention_tiled_3d_lower_llvm"); } + +std::string t3d942_lower(const py::dict& d, const std::string& arch) +{ + require_arch(arch); + Store st; + rocke_unified_attention_3d_tiled_spec_t s = t3d_build(d, st); + const char* target = arch.c_str(); + std::string out = t3d942_segment_ll(d, s, target); + rocke_unified_attention_reduce_tiled_spec_t r + = rocke_unified_attention_reduce_tiled_spec_default(); + r.head_size = s.head_size; + r.num_query_heads = s.num_query_heads; + r.num_kv_heads = s.num_kv_heads; + r.dtype = s.dtype; + r.num_segments = s.num_segments; + r.has_waves_per_eu = s.has_waves_per_eu; + r.waves_per_eu = s.waves_per_eu; + r.emit_lse = a_bool(d, "emit_lse", r.emit_lse); + char* ll = nullptr; + char err[ROCKE_ERR_MSG_CAP]; + err[0] = '\0'; + rocke_status_t status = rocke_build_unified_attention_reduce_tiled_gfx942_lower_to_llvm( + &r, target, resolve_python_llvm_flavor(), &ll, err, sizeof err); + out += take_ll(status, ll, err, "rocke_engine.gfx942_attention_tiled_3d_lower_llvm"); + return out; +} std::string t3d942_serialize(const py::dict& d, const std::string& arch) { require_arch(arch); @@ -1231,6 +1303,7 @@ std::string t3d950_lower(const py::dict& d, const std::string& arch) r.num_segments = s.num_segments; r.has_waves_per_eu = s.has_waves_per_eu; r.waves_per_eu = s.waves_per_eu; + r.emit_lse = a_bool(d, "emit_lse", r.emit_lse); rocke_ir_builder_t b; rocke_ir_builder_init(&b, "attention_tiled_3d_reduce"); rocke_kernel_def_t* k = rocke_build_unified_attention_reduce_tiled_gfx950(&b, &r, a); @@ -1381,6 +1454,7 @@ static py::dict spec_to_dict(const rocke_attention_tiled_2d_spec_t& s) d["use_i64_kv_addr"] = s.use_i64_kv_addr; d["use_k_single_buffer"] = s.use_k_single_buffer; d["use_fp8_mfma_qk"] = s.use_fp8_mfma_qk; + d["emit_lse"] = s.emit_lse; return d; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/abi.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/abi.h index 4c0e3d2a5752..a6ea37b79927 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/abi.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/abi.h @@ -86,8 +86,11 @@ extern "C" { * 4: rocke_memory_caps_t gained has_global_float_atomic_add at its end, which * changes its size and with it the offset of rocke_arch_target_t.limits. * 5: rocke_tensor_view_t gained buffer (rocke_tv_buffer_resource_t) at its end, - * which changes its size and that of every struct embedding a view. */ -#define ROCKE_ABI_VERSION 5 + * which changes its size and that of every struct embedding a view. + * 6: rocke_attention_tiled_2d_spec_t (gfx942) and + * rocke_unified_attention_reduce_tiled_spec_t (gfx942, gfx950) gained + * emit_lse at their end, which changes their size. */ +#define ROCKE_ABI_VERSION 6 /* Wire ABI: the CBOR recipe/bundle/guard vocabulary this engine understands. * 2: list-tagged ("l") kernel attrs such as cluster_dims; level-1 engines drop them. diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h index 1014518abdae..8b432fb14cc1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h @@ -179,6 +179,10 @@ typedef struct rocke_attention_tiled_2d_spec bool use_global_load_lds_k; /* False */ bool use_q_major_grid; /* False */ bool use_qblock_reverse; /* False (gfx950: longest-first query blocks) */ + /* Also write the natural-log softmax normalizer (LSE) as FP32 + * [total_q, num_query_heads] through a trailing lse_ptr param. Off emits + * no new IR; on appends the "lse" kernel-name token. */ + bool emit_lse; /* False (Python kw_only) */ } rocke_attention_tiled_2d_spec_t; /* Materialise every defaulted field (the dataclass defaults). The required diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h index 344729922718..55c8ebbeb4a2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h @@ -178,6 +178,8 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* num_seqs_p; rocke_value_t* bt_stride_p; rocke_value_t* qq_bias_stride0_p; + /* Trailing lse_ptr (F32*), declared only when spec->emit_lse; else NULL. */ + rocke_value_t* lse; /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ rocke_value_t* q_block_global_idx; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h index 6b689a020685..6729eff640d4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h @@ -224,6 +224,9 @@ typedef struct rocke_unified_attention_reduce_tiled_spec bool has_waves_per_eu; /* Optional[int] None */ int waves_per_eu; + /* emit_lse (False): add a trailing lse_ptr param and write the FP32 natural-log + * LSE [total_q, num_query_heads, 1]; kernel name gains a trailing "lse". */ + bool emit_lse; } rocke_unified_attention_reduce_tiled_spec_t; rocke_unified_attention_reduce_tiled_spec_t rocke_unified_attention_reduce_tiled_spec_default(void); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h index 67dd78e8786e..1f58dc76c378 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h @@ -321,6 +321,7 @@ typedef struct rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* seg_max; /* F32* readonly */ rocke_value_t* seg_l; /* F32* readonly (segm_expsum) */ rocke_value_t* red_seq_lens; /* I32* readonly (_seq_lens, unused body) */ + rocke_value_t* red_lse; /* F32* noalias writeonly (lse_ptr); NULL unless emit_lse */ /* ---------- grid ids (lines 1032-1034) ---------- */ rocke_value_t* q_token; /* block_id_x() */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h index dc945f0f4579..f6137622d663 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h @@ -175,6 +175,7 @@ typedef struct rocke_gfx950_attn2d_build_ctx rocke_value_t* num_seqs_p; rocke_value_t* bt_stride_p; rocke_value_t* qq_bias_stride0_p; + rocke_value_t* lse; /* lse_ptr (spec->emit_lse only, else NULL) */ /* ---- grid ids + wave decomposition ---- */ rocke_value_t* q_block_global_idx; @@ -324,6 +325,9 @@ typedef struct rocke_gfx950_attn2d_build_ctx rocke_value_t* hoist_q_pos[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; rocke_value_t* hoist_q_head[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; rocke_value_t* hoist_row_mask[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; + /* Python hoist_qh_r: the query head, never replaced by the ALiBi slope + * (hoist_q_head is). Read by the emit_lse epilogue stores. */ + rocke_value_t* hoist_qh_r[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; /* ---- transposed-32x32 invariant-hoist state (Python 2446-2480) ---- * * Computed once in the LICM hoist when TRANSPOSED_INVARIANT_HOIST is set; @@ -352,6 +356,7 @@ typedef struct rocke_gfx950_attn2d_build_ctx /* ---- final results read by the epilogue (lines 3585-3818) ---- */ rocke_value_t* l_final[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; + rocke_value_t* m_final[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; /* emit_lse only */ rocke_value_t* acc_final[ROCKE_GFX950_ATTN2D_MAX_ACCS]; rocke_value_t* rcp_l[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; rocke_value_t* l_zero[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h index 6506e70467d6..b674a5c2b2a3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h @@ -220,6 +220,9 @@ typedef struct rocke_unified_attention_reduce_tiled_spec bool has_waves_per_eu; /* Optional[int] None */ int waves_per_eu; + /* emit_lse (False): add a trailing lse_ptr param and write the FP32 natural-log + * LSE [total_q, num_query_heads, 1]; kernel name gains a trailing "lse". */ + bool emit_lse; } rocke_unified_attention_reduce_tiled_spec_t; rocke_unified_attention_reduce_tiled_spec_t rocke_unified_attention_reduce_tiled_spec_default(void); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h index 4363fedd6162..69228aa7003a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h @@ -361,6 +361,7 @@ typedef struct rocke_gfx950_attention_tiled_3d_build_ctx rocke_value_t* seg_max; /* F32* readonly */ rocke_value_t* seg_l; /* F32* readonly (segm_expsum) */ rocke_value_t* red_seq_lens; /* I32* readonly (_seq_lens, unused body) */ + rocke_value_t* red_lse; /* F32* noalias writeonly (lse_ptr); NULL unless emit_lse */ /* ---------- grid ids (lines 1029-1031) ---------- */ rocke_value_t* q_token; /* block_id_x() */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_lse_store_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_lse_store_internal.h new file mode 100644 index 000000000000..0c62927785ab --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_lse_store_internal.h @@ -0,0 +1,60 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* LSE store internals. Mirrors library/kernels/common/_lse_store.py + * (lse_store, lse_row_has_mass) and natural_lse_from_log2_stats from + * rocke/helpers/attention_fwd_ext.py, op for op and in the same order, so the + * two engines emit byte-identical IR. Every builder call is its own statement: + * C leaves the evaluation order of call arguments unspecified, and the order + * the ops are created in decides their SSA names. Not a public ABI. */ +#ifndef ROCKE_INSTANCE_LSE_STORE_INTERNAL_H +#define ROCKE_INSTANCE_LSE_STORE_INTERNAL_H + +#include + +#include "rocke/ir.h" + +/* LN2 in attention_fwd_ext.py. */ +#define ROCKE_LSE_LN2 0.6931471805599453 + +/* natural_lse_from_log2_stats: LN2 * (m_log2 + log2(row_sum)). */ +static inline rocke_value_t* rocke_natural_lse_from_log2_stats(rocke_ir_builder_t* b, + rocke_value_t* m_log2, + rocke_value_t* row_sum) +{ + rocke_value_t* lg = rocke_b_log2(b, row_sum); + rocke_value_t* sum = rocke_b_fadd(b, m_log2, lg); + rocke_value_t* ln2 = rocke_b_const_f32(b, ROCKE_LSE_LN2); + return rocke_b_fmul(b, sum, ln2); +} + +/* lse_row_has_mass: !(l == 0) with NaN counted as mass, the exact inverse of + * the tiled 2D O guard fcmp("oeq", l, 0). Built as one OR uno, not + * lnot(oeq): the HIP lowerer lowers i1 NOT to ~bool, which is always true. */ +static inline rocke_value_t* rocke_lse_row_has_mass(rocke_ir_builder_t* b, rocke_value_t* l) +{ + rocke_value_t* zero = rocke_b_const_f32(b, 0.0); + rocke_value_t* ne = rocke_b_fcmp(b, "one", l, zero); + rocke_value_t* nan = rocke_b_fcmp(b, "uno", l, zero); + return rocke_b_lor(b, ne, nan); +} + +/* lse_store: one FP32 store of the natural-log LSE; -inf when the row has no + * mass. The caller guards row existence and writer ownership. */ +static inline void rocke_lse_store(rocke_ir_builder_t* b, + rocke_value_t* dst, + rocke_value_t* element_offset, + rocke_value_t* m_log2, + rocke_value_t* l, + rocke_value_t* row_has_mass) +{ + rocke_value_t* one = rocke_b_const_f32(b, 1.0); + rocke_value_t* safe_l = rocke_b_select(b, row_has_mass, l, one); + rocke_value_t* zero = rocke_b_const_f32(b, 0.0); + rocke_value_t* safe_m = rocke_b_select(b, row_has_mass, m_log2, zero); + rocke_value_t* value = rocke_natural_lse_from_log2_stats(b, safe_m, safe_l); + rocke_value_t* neg_inf = rocke_b_const_f32(b, -INFINITY); + value = rocke_b_select(b, row_has_mass, value, neg_inf); + rocke_b_global_store(b, dst, element_offset, value, 4); +} + +#endif /* ROCKE_INSTANCE_LSE_STORE_INTERNAL_H */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d.cpp index 821229720b5d..427151670a77 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d.cpp @@ -129,6 +129,7 @@ rocke_attention_tiled_2d_spec_t rocke_attention_tiled_2d_spec_default(void) s.use_global_load_lds_k = false; s.use_q_major_grid = false; s.use_qblock_reverse = false; + s.emit_lse = false; return s; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp index 6e9f5ca52c09..d6cdd99b88f5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp @@ -398,6 +398,18 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + /* Appended LAST and only when requested, so the LSE-off ABI is unchanged. */ + if(spec->emit_lse) + { + memset(&o, 0, sizeof(o)); + o.noalias = true; + o.noalias_set = true; + o.writeonly = true; + o.writeonly_set = true; + o.align = 4; + o.align_set = true; + ctx->lse = rocke_b_param(b, "lse_ptr", ptr_f32, &o); + } } /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp index c1f457fc2545..cf991e78e03d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp @@ -18,6 +18,7 @@ #include "rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h" /* rocke__mfma_32x32_c_row */ #include "rocke/helper_rocke.helpers.transforms.h" /* descriptor offset */ #include "rocke/instance_gfx942_attention_tiled_2d_internal.h" +#include "rocke/instance_lse_store_internal.h" /* rocke_lse_store / _row_has_mass */ /* q_desc.offset(b, token=tok, head=hd, dim=dimv) -> offset (the `_` valid is * discarded by the Python epilogue). Mirrors TensorDescriptor.offset(**upper). @@ -40,6 +41,43 @@ static rocke_value_t* epilogue_q_desc_offset(rocke_gfx942_attn2d_build_ctx_t* ct return offset; } +/* Python _store_row_lse(row, slot, writer): write LSE for CTA row ``row`` from + * softmax state ``slot``, from the one ``writer`` lane among the lanes holding + * a replicated copy of the row's (m, l). Same token/head/validity mapping as + * the O stores. Each builder call is its own statement, in Python order. */ +static void epilogue_store_row_lse(rocke_gfx942_attn2d_build_ctx_t* ctx, + rocke_value_t* row, + rocke_value_t* m_final, + rocke_value_t* l_final, + rocke_value_t* writer) +{ + rocke_ir_builder_t* b = ctx->b; + rocke_value_t* pos_div = rocke_b_div(b, row, rocke_b_const_i32(b, ctx->NQK)); + rocke_value_t* pos = rocke_b_add(b, ctx->qb_start_pos, pos_div); + rocke_value_t* qh_mul = rocke_b_mul(b, ctx->kv_head_idx, rocke_b_const_i32(b, ctx->NQK)); + rocke_value_t* qh_mod = rocke_b_mod(b, row, rocke_b_const_i32(b, ctx->NQK)); + rocke_value_t* qh = rocke_b_add(b, qh_mul, qh_mod); + rocke_value_t* ok_pos = rocke_b_cmp_lt(b, pos, ctx->cur_batch_q_len); + rocke_value_t* ok_qh = rocke_b_cmp_lt(b, qh, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* ok = rocke_b_land(b, ok_pos, ok_qh); + if(ctx->VALID_ROWS < ctx->BLOCK_M) + { + rocke_value_t* ok_vr = rocke_b_cmp_lt(b, row, rocke_b_const_i32(b, ctx->VALID_ROWS)); + ok = rocke_b_land(b, ok, ok_vr); + } + rocke_value_t* cond = rocke_b_land(b, writer, ok); + rocke_if_t iff = rocke_b_scf_if(b, cond); + rocke_b_region_enter(b, iff.then_region); + { + rocke_value_t* tok = rocke_b_add(b, ctx->cu_q_start, pos); + rocke_value_t* tok_mul = rocke_b_mul(b, tok, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* offset = rocke_b_add(b, tok_mul, qh); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, l_final); + rocke_lse_store(b, ctx->lse, offset, m_final, l_final, has_mass); + } + rocke_b_region_leave(b); +} + rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; @@ -87,6 +125,11 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ /* l_final[r] = final[2*r + 1] */ for(r = 0; r < SOFTMAX_STATE_SLOTS; ++r) ctx->l_final[r] = final[2 * r + 1]; + /* m_final[r] = final[2*r] (log2 units); read only for LSE. */ + rocke_value_t* m_final[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE] = {}; + if(ctx->spec->emit_lse) + for(r = 0; r < SOFTMAX_STATE_SLOTS; ++r) + m_final[r] = final[2 * r]; /* acc_final indexed by (n * ACC_M_ATOMS + atom). */ for(n = 0; n < ACC_N_TILES; ++n) @@ -152,6 +195,26 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* inv_l_t = rocke_b_rcp(b, ctx->l_final[0]); rocke_value_t* l_zero_t = rocke_b_fcmp(b, "oeq", ctx->l_final[0], zero_f); + if(ctx->spec->emit_lse) + { + /* One query row per lane (lane % 32); lanes l and l^32 hold the + * same state. The lower half (lane_half32 == 0) writes slot 0. */ + rocke_value_t* lane_half32 = (ctx->lane_half32_v != NULL) + ? ctx->lane_half32_v + : rocke_b_div(b, lane, rocke_b_const_i32(b, 32)); + rocke_value_t* lo_half = rocke_b_cmp_eq(b, lane_half32, rocke_b_const_i32(b, 0)); + rocke_value_t* lse_cond = rocke_b_land(b, op_mask_t, lo_half); + rocke_if_t lse_if = rocke_b_scf_if(b, lse_cond); + rocke_b_region_enter(b, lse_if.then_region); + { + rocke_value_t* tok = rocke_b_add(b, cu_q_start, op_pos_t); + rocke_value_t* tok_mul = rocke_b_mul(b, tok, rocke_b_const_i32(b, NUM_QH)); + rocke_value_t* offset = rocke_b_add(b, tok_mul, op_qh_t); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->l_final[0]); + rocke_lse_store(b, ctx->lse, offset, m_final[0], ctx->l_final[0], has_mass); + } + rocke_b_region_leave(b); + } for(n = 0; n < ACC_N_TILES; ++n) { @@ -225,6 +288,20 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* out_token32 = rocke_b_add(b, cu_q_start, op_pos32_base); rocke_value_t* out_base32_base = epilogue_q_desc_offset(ctx, out_token32, op_qh32_base, rocke_b_const_i32(b, 0)); + if(ctx->spec->emit_lse) + { + /* Slot r is C-row wave_row_base + _mfma_32x32_c_row(lane, r); lanes + * 0 and 32 (lane_col32 == 0) write. Before the first Acc_lds sync. */ + rocke_value_t* lane_col32_w + = (ctx->lane_col32_q32_v != NULL) ? ctx->lane_col32_q32_v : lane_col32; + rocke_value_t* writer32 = rocke_b_cmp_eq(b, lane_col32_w, rocke_b_const_i32(b, 0)); + for(r = 0; r < SOFTMAX_STATE_SLOTS; ++r) + { + rocke_value_t* c_row = rocke__mfma_32x32_c_row(b, lane, r); + rocke_value_t* lse_row = rocke_b_add(b, wave_row_base, c_row); + epilogue_store_row_lse(ctx, lse_row, m_final[r], ctx->l_final[r], writer32); + } + } for(n = 0; n < ACC_N_TILES; ++n) { @@ -377,5 +454,22 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_b_sync(b); } + if(ctx->spec->emit_lse) + { + /* After the O stripes, not before: emitted ahead of them, the LSE store + * let LLVM fuse the O multiply and FP16 convert into v_fma_mixlo_f16 + * (one rounding; LSE-off rounds twice), moving FP16 O by 1 ulp versus + * LSE-off on MI300X/MI355X (ROCm 10). Here the O ISA matches LSE-off. + * Slot r is row wave_row_base + _in_warp_row(r), shared by the 16 + * lane_col lanes; lane_col == 0 (lanes 0/16/32/48) writes. */ + rocke_value_t* writer16 = rocke_b_cmp_eq(b, lane_col, rocke_b_const_i32(b, 0)); + for(r = 0; r < SOFTMAX_STATE_SLOTS; ++r) + { + rocke_value_t* iw_row = rocke_gfx942_attn2d_in_warp_row(ctx, r); + rocke_value_t* lse_row = rocke_b_add(b, wave_row_base, iw_row); + epilogue_store_row_lse(ctx, lse_row, m_final[r], ctx->l_final[r], writer16); + } + } + return rocke_ir_builder_kernel(b); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index 17a5653dc8a1..25da0bf46b41 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -623,6 +623,7 @@ static rocke_status_t s->use_k_sliced_ldsseq ? "ldsseq" : "", s->use_iglp_opt ? "iglp1" : "", s->use_k_single_buffer ? "k1buf" : "", + s->emit_lse ? "lse" : "", }; size_t num_parts = sizeof(parts) / sizeof(parts[0]); return rocke_kernel_name_join( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index 504e0397492d..e12a8548a223 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -392,6 +392,7 @@ rocke_unified_attention_reduce_tiled_spec_t rocke_unified_attention_reduce_tiled memset(&s, 0, sizeof(s)); s.has_waves_per_eu = false; s.waves_per_eu = 0; + s.emit_lse = false; return s; } @@ -409,7 +410,7 @@ int rocke_unified_attention_reduce_tiled_spec_kernel_name( const rocke_unified_attention_reduce_tiled_spec_t* s, char* buf, size_t cap) { /* kernel_name_join("rocke_uattn_reduce_tiled_gfx942", d, h, - * seg, dtype) */ + * seg, dtype[, "lse"]) */ char d_part[32]; char h_part[32]; char seg_part[32]; @@ -432,6 +433,10 @@ int rocke_unified_attention_reduce_tiled_spec_kernel_name( parts[np++] = h_part; parts[np++] = seg_part; parts[np++] = (s->dtype != NULL) ? s->dtype : ""; + if(s->emit_lse) + { + parts[np++] = "lse"; + } st = rocke_kernel_name_join("", parts, np, NULL, NULL, 0, buf, cap, &out_len); if(st != ROCKE_OK) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp index b810032b2ec0..ee2583d176dd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp @@ -29,6 +29,7 @@ #include #include "rocke/instance_gfx942_attention_tiled_3d_internal.h" +#include "rocke/instance_lse_store_internal.h" /* ============================================================ * * Per-slot caches shared between reduce_max_pass and reduce_combine_pass. @@ -188,6 +189,20 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( ctx->red_seq_lens = rocke_b_param(b, "seq_lens_ptr", rocke_ptr_type(b, rocke_i32(), "global"), &opts); + /* lse_ptr : ptr, noalias, writeonly, align=4 (emit_lse only) */ + ctx->red_lse = NULL; + if(spec != NULL && spec->emit_lse) + { + memset(&opts, 0, sizeof(opts)); + opts.noalias = true; + opts.noalias_set = true; + opts.writeonly = true; + opts.writeonly_set = true; + opts.align = 4; + opts.align_set = true; + ctx->red_lse = rocke_b_param(b, "lse_ptr", rocke_ptr_type(b, rocke_f32(), "global"), &opts); + } + /* ---- grid ids + thread (lines 1032-1034) ---- */ ctx->q_token = rocke_b_block_id_x(b); ctx->q_head = rocke_b_block_id_y(b); @@ -508,4 +523,26 @@ void rocke_gfx942_attention_tiled_3d_reduce_normalize_pass( out_idx = rocke__red_out_offset(ctx, ctx->q_token, ctx->q_head, d); rocke_b_global_store(b, ctx->out, out_idx, scalar_out, 2); } + + /* emit_lse: every lane holds the same butterfly-reduced overall_max/expsum, + * so lane 0 writes. overall_max is in log2 units; lse_store converts to + * natural log and writes -inf for a row with no mass (sum 0); a NaN sum + * keeps NaN, as O does (!(l == 0), never l > 0). One builder call per + * statement, in Python order. */ + if(ctx->reduce_spec != NULL && ctx->reduce_spec->emit_lse) + { + rocke_value_t* c0 = rocke_b_const_i32(b, 0); + rocke_value_t* is_writer = rocke_b_cmp_eq(b, ctx->tid, c0); + rocke_if_t guard = rocke_b_scf_if(b, is_writer); + rocke_b_region_enter(b, guard.then_region); + { + rocke_value_t* nqh = rocke_b_const_i32(b, cfg->NUM_QH); + rocke_value_t* row = rocke_b_mul(b, ctx->q_token, nqh); + rocke_value_t* lse_index = rocke_b_add(b, row, ctx->q_head); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->overall_expsum); + rocke_lse_store( + b, ctx->red_lse, lse_index, ctx->overall_max, ctx->overall_expsum, has_mass); + } + rocke_b_region_leave(b); + } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp index c6c4ab018e20..5a1aa6edb1d6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp @@ -292,6 +292,17 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + if(spec->emit_lse) + { + memset(&o, 0, sizeof(o)); + o.noalias = true; + o.noalias_set = true; + o.writeonly = true; + o.writeonly_set = true; + o.align = 4; + o.align_set = true; + ctx->lse = rocke_b_param(b, "lse_ptr", ptr_f32, &o); + } } /* ---- grid ids + wave decomposition (lines 891-904) ---- * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp index dabd67225c4a..0916cf1d0783 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp @@ -98,6 +98,7 @@ static rocke_status_t rocke__fastkv_gfx950_base_name(const rocke_attention_tiled s->use_fp8_mfma_qk ? "fp8mfma" : "", s->use_fp8_mfma_pv ? "fp8pv" : "", s->use_register_pv ? "regpv" : "", + s->emit_lse ? "lse" : "", }; size_t num_parts = sizeof(parts) / sizeof(parts[0]); return rocke_kernel_name_join( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp index 21d361f81eca..2a16c0c13624 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp @@ -36,6 +36,7 @@ #include "rocke/helper_rocke.helpers.mfma_attention.h" /* 32x32x16_for_dtype */ #include "rocke/instance_gfx950_attention_tiled_2d.h" /* mfma_32x32_c_row/_col */ #include "rocke/instance_gfx950_attention_tiled_2d_internal.h" +#include "rocke/instance_lse_store_internal.h" /* rocke_lse_store / row_has_mass */ #include @@ -751,6 +752,10 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ * The driver stashed kvloop.op->results into ctx->out_carry. */ for(r = 0; r < ctx->SOFTMAX_STATE_SLOTS; ++r) ctx->l_final[r] = ctx->out_carry[2 * r + 1]; + /* emit_lse: m is in log2 units (scores carry scale*log2e); no IR. */ + if(ctx->spec->emit_lse) + for(r = 0; r < ctx->SOFTMAX_STATE_SLOTS; ++r) + ctx->m_final[r] = ctx->out_carry[2 * r]; for(n = 0; n < ctx->ACC_N_TILES; ++n) for(atom = 0; atom < ctx->ACC_M_ATOMS; ++atom) ctx->acc_final[n * ctx->ACC_M_ATOMS + atom] @@ -812,6 +817,25 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ inv_l_t = rocke_b_rcp(b, ctx->l_final[0]); l_zero_t = rocke_b_fcmp(b, "oeq", ctx->l_final[0], ctx->zero_f); + if(ctx->spec->emit_lse) + { + /* One row per lane; lanes l and l^32 hold the same (m, l), so only + * the low half writes. Emitted once, outside the n-loop. */ + rocke_value_t* zero_i = rocke_b_const_i32(b, 0); + rocke_value_t* low_half = rocke_b_cmp_eq(b, ctx->lane_half32_v, zero_i); + rocke_value_t* writer = rocke_b_land(b, op_mask_t, low_half); + rocke_if_t lse_if = rocke_b_scf_if(b, writer); + rocke_b_region_enter(b, lse_if.then_region); + { + rocke_value_t* tok = rocke_b_add(b, ctx->cu_q_start, op_pos_t); + rocke_value_t* nqh = rocke_b_const_i32(b, ctx->NUM_QH); + rocke_value_t* tok_qh = rocke_b_mul(b, tok, nqh); + rocke_value_t* off = rocke_b_add(b, tok_qh, op_qh_t); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->l_final[0]); + rocke_lse_store(b, ctx->lse, off, ctx->m_final[0], ctx->l_final[0], has_mass); + } + rocke_b_region_leave(b); + } for(n = 0; n < ctx->ACC_N_TILES; ++n) { rocke_value_t* acc32 = rocke_gfx950_attn2d_acc_final_get(ctx, n, 0); @@ -900,6 +924,29 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ if(!rocke_transforms_descriptor_offset( b, ctx->q_desc, in_names, in_values, 3, &out_base32_base, NULL)) return NULL; + if(ctx->spec->emit_lse) + { + /* Slot r's row is wave_row_base + _mfma_32x32_c_row(lane, r); all 32 + * lanes of a half agree, lane_col32 == 0 (lanes 0 and 32) writes. */ + rocke_value_t* zero_i = rocke_b_const_i32(b, 0); + rocke_value_t* writer32 = rocke_b_cmp_eq(b, ctx->lane_col32_v, zero_i); + for(r = 0; r < ctx->REGS_PER_LANE; ++r) + { + rocke_value_t* guard = rocke_b_land(b, writer32, ctx->hoist_row_mask[r]); + rocke_if_t lse_if = rocke_b_scf_if(b, guard); + rocke_b_region_enter(b, lse_if.then_region); + { + rocke_value_t* tok = rocke_b_add(b, ctx->cu_q_start, ctx->hoist_q_pos[r]); + rocke_value_t* nqh = rocke_b_const_i32(b, ctx->NUM_QH); + rocke_value_t* tok_qh = rocke_b_mul(b, tok, nqh); + rocke_value_t* off = rocke_b_add(b, tok_qh, ctx->hoist_qh_r[r]); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->l_final[r]); + rocke_lse_store( + b, ctx->lse, off, ctx->m_final[r], ctx->l_final[r], has_mass); + } + rocke_b_region_leave(b); + } + } for(n = 0; n < ctx->ACC_N_TILES; ++n) { @@ -1018,7 +1065,6 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ if(!rocke_transforms_descriptor_offset( b, ctx->q_desc, in_names, in_values, 3, &out_base, NULL)) return NULL; - for(stripe = 0; stripe < ctx->OUT_STRIPES; ++stripe) { int n_start = stripe * N_TILES_PER_STRIPE; @@ -1075,6 +1121,33 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ if(stripe + 1 < ctx->OUT_STRIPES) rocke_b_sync(b); } + if(ctx->spec->emit_lse) + { + /* After the O stripes, not before: emitted ahead of them, the LSE + * store let LLVM fuse the O multiply and FP16 convert into + * v_fma_mixlo_f16 (one rounding; LSE-off rounds twice), moving FP16 + * O by 1 ulp versus LSE-off on MI300X/MI355X (ROCm 10). Here the O + * ISA matches LSE-off. Slot r's row is wave_row_base + + * _in_warp_row(r); the 16 lanes of a row group agree, lane_col == 0 + * (one per lane_rg) writes. */ + rocke_value_t* zero_i = rocke_b_const_i32(b, 0); + rocke_value_t* writer16 = rocke_b_cmp_eq(b, ctx->lane_col_v, zero_i); + for(r = 0; r < ctx->REGS_PER_LANE; ++r) + { + rocke_value_t* guard = rocke_b_land(b, writer16, ctx->hoist_row_mask[r]); + rocke_if_t lse_if = rocke_b_scf_if(b, guard); + rocke_b_region_enter(b, lse_if.then_region); + { + rocke_value_t* tok = rocke_b_add(b, ctx->cu_q_start, ctx->hoist_q_pos[r]); + rocke_value_t* nqh = rocke_b_const_i32(b, ctx->NUM_QH); + rocke_value_t* tok_qh = rocke_b_mul(b, tok, nqh); + rocke_value_t* off = rocke_b_add(b, tok_qh, ctx->hoist_qh_r[r]); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->l_final[r]); + rocke_lse_store(b, ctx->lse, off, ctx->m_final[r], ctx->l_final[r], has_mass); + } + rocke_b_region_leave(b); + } + } } (void)i; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp index 6a3fd1130bf7..e9d2c146e699 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp @@ -261,6 +261,7 @@ void rocke_gfx950_attn2d_emit_licm_hoist(rocke_gfx950_attn2d_build_ctx_t* ctx) ctx->hoist_in_warp_row[reg] = row; ctx->hoist_q_pos[reg] = qp_r; ctx->hoist_q_head[reg] = qh_r; + ctx->hoist_qh_r[reg] = qh_r; ctx->hoist_row_mask[reg] = row_ok; ctx->hoist_state_row[reg] = causal_lim; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp index 6132f8356e2a..b2c064644323 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp @@ -743,6 +743,19 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + /* Trailing LSE output (Python: after the optional additive-bias params, + * which this twin does not declare). Off -> no param, ABI unchanged. */ + if(spec->emit_lse) + { + memset(&o, 0, sizeof(o)); + o.noalias = true; + o.noalias_set = true; + o.writeonly = true; + o.writeonly_set = true; + o.align = 4; + o.align_set = true; + ctx->lse = rocke_b_param(b, "lse_ptr", ptr_f32, &o); + } } /* ---- grid ids + wave decomposition (Py891-904) ---- */ @@ -1127,6 +1140,7 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ s->use_fp8_mfma_pv ? "fp8pv" : "", s->use_register_pv ? "regpv" : "", s->use_qblock_reverse ? "qrev" : "", + s->emit_lse ? "lse" : "", }; size_t num_parts = sizeof(parts) / sizeof(parts[0]); return rocke_kernel_name_join( diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp index 8b22ac7b10db..939d46c55b52 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp @@ -221,6 +221,7 @@ rocke_unified_attention_reduce_tiled_spec_t rocke_unified_attention_reduce_tiled memset(&s, 0, sizeof(s)); s.has_waves_per_eu = false; s.waves_per_eu = 0; + s.emit_lse = false; return s; } @@ -517,7 +518,7 @@ int rocke_gfx950_unified_attention_reduce_tiled_spec_kernel_name( const rocke_unified_attention_reduce_tiled_spec_t* s, char* buf, size_t cap) { /* kernel_name_join("rocke_uattn_reduce_tiled", d, h, - * seg, dtype) */ + * seg, dtype[, "lse"]) */ char d_part[32]; char h_part[32]; char seg_part[32]; @@ -540,6 +541,10 @@ int rocke_gfx950_unified_attention_reduce_tiled_spec_kernel_name( parts[np++] = h_part; parts[np++] = seg_part; parts[np++] = (s->dtype != NULL) ? s->dtype : ""; + if(s->emit_lse) + { + parts[np++] = "lse"; + } st = rocke_kernel_name_join("", parts, np, NULL, NULL, 0, buf, cap, &out_len); if(st != ROCKE_OK) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp index 0ea23c18ea85..91f4a05f1728 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp @@ -28,6 +28,7 @@ #include #include "rocke/instance_gfx950_attention_tiled_3d_internal.h" +#include "rocke/instance_lse_store_internal.h" /* ============================================================ * * Per-slot caches shared between reduce_max_pass and reduce_combine_pass. @@ -189,6 +190,20 @@ void rocke_gfx950_attention_tiled_3d_reduce_declare_and_prologue( ctx->red_seq_lens = rocke_b_param(b, "seq_lens_ptr", rocke_ptr_type(b, rocke_i32(), "global"), &opts); + /* lse_ptr : ptr, noalias, writeonly, align=4 (emit_lse only) */ + ctx->red_lse = NULL; + if(spec != NULL && spec->emit_lse) + { + memset(&opts, 0, sizeof(opts)); + opts.noalias = true; + opts.noalias_set = true; + opts.writeonly = true; + opts.writeonly_set = true; + opts.align = 4; + opts.align_set = true; + ctx->red_lse = rocke_b_param(b, "lse_ptr", rocke_ptr_type(b, rocke_f32(), "global"), &opts); + } + /* ---- grid ids + thread (lines 1029-1031) ---- */ ctx->q_token = rocke_b_block_id_x(b); ctx->q_head = rocke_b_block_id_y(b); @@ -507,4 +522,26 @@ void rocke_gfx950_attention_tiled_3d_reduce_normalize_pass( out_idx = rocke__red_out_offset(ctx, ctx->q_token, ctx->q_head, d); rocke_b_global_store(b, ctx->out, out_idx, scalar_out, 2); } + + /* emit_lse: every lane holds the same butterfly-reduced overall_max/expsum, + * so lane 0 writes. overall_max is in log2 units; lse_store converts to + * natural log and writes -inf for a row with no mass (sum 0); a NaN sum + * keeps NaN, as O does (!(l == 0), never l > 0). One builder call per + * statement, in Python order. */ + if(ctx->reduce_spec != NULL && ctx->reduce_spec->emit_lse) + { + rocke_value_t* c0 = rocke_b_const_i32(b, 0); + rocke_value_t* is_writer = rocke_b_cmp_eq(b, ctx->tid, c0); + rocke_if_t guard = rocke_b_scf_if(b, is_writer); + rocke_b_region_enter(b, guard.then_region); + { + rocke_value_t* nqh = rocke_b_const_i32(b, cfg->NUM_QH); + rocke_value_t* row = rocke_b_mul(b, ctx->q_token, nqh); + rocke_value_t* lse_index = rocke_b_add(b, row, ctx->q_head); + rocke_value_t* has_mass = rocke_lse_row_has_mass(b, ctx->overall_expsum); + rocke_lse_store( + b, ctx->red_lse, lse_index, ctx->overall_max, ctx->overall_expsum, has_mass); + } + rocke_b_region_leave(b); + } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/known_gaps.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/known_gaps.md index 4258674b4674..33afb38dc22a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/known_gaps.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/development/known_gaps.md @@ -244,40 +244,42 @@ mirrors; until then, keep emission changes gated by the two tests above. --- -## 7. Unified split-KV reduce LSE output (`emit_lse`) — Python only - -**Subsystem:** `instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp`, -`instances/gfx950/attention_tiled_3d_reduce.cpp`, -`include/rocke/instance_gfx9{42,50}_attention_tiled_3d.h`, -`instances/common/attention_unified_selectors.cpp` (and its header) - -`UnifiedAttentionReduceTiledSpec.emit_lse` (both arches, -`library/kernels/gfx9{42,50}/attention_tiled_3d.py`) makes the split-KV reduce -kernel take a trailing `lse_ptr` and write one FP32 natural-log LSE per -(query token, query head) from lane 0 via `kernels/common/_lse_store.py`. The C -engine does not mirror it. With `emit_lse=False` (the default) the reduce gets -no `lse_ptr` and no store, so this gap adds nothing to an LSE-off kernel; and no -C entry point can request LSE, so no C-built kernel reaches the missing path. -This entry covers `emit_lse` only, not other changes to the reduce kernel. +## 7. Unified attention LSE output (`emit_lse`) — not requestable from C + +**Subsystem:** `instances/common/attention_unified_selectors.cpp` (and its +header), `rocke_unified_attn_problem_t` + +The C engine builds every LSE-on unified attention kernel byte-identically to +Python: the split-KV reduce (gfx942, gfx950) and the default tiled 2D builders +(gfx942, gfx950, including the gfx950 `fastkv_regp` wrapper). They share +`include/rocke/instance_lse_store_internal.h`, the C twin of +`kernels/common/_lse_store.py`. The parity emitters build LSE-on configs for +each store site (`gfx9{42,50}_attention_tiled_2d_emit`, the gfx950 +`_fastkv_regp_emit`, `gfx9{42,50}_attention_tiled_3d_emit`), so +`check_byte_identity` gates them. + +What is still Python-only is *asking* for LSE: the request side. | What's missing | Where | |---|---| -| `emit_lse` field (default 0) and the `lse` kernel-name suffix | `rocke_unified_attention_reduce_tiled_spec_t` and `*_reduce_tiled_spec_kernel_name()`, both arches | -| Trailing `lse_ptr` param and the `tid == 0` LSE store | both C reduce builders | -| A C twin of `_lse_store.py` | `helpers/`; none exists (its dense-attention users have no C twin either) | -| Gate coverage of an LSE-on reduce | `library/tests/parity/gfx9{42,50}_attention_tiled_3d_emit.{py,c}` emit the LSE-off reduce only; neither takes `emit_lse` | -| `emit_lse` on the unified problem and the `_reject_lse` support gate (`library/kernels/common/attention_unified.py`) | `rocke_unified_attn_problem_t`; the C selectors carry no support gates to mirror it into. No codegen selector reads `emit_lse`: the selected split-KV spec differs only in `reduce_spec.emit_lse` between LSE on and off, so C selector output is unaffected | - -**Why deferred:** owner decision for ROCm/rocm-libraries#13955 (time-sensitive; -the C engine is expected to be deprecated). Python is the only engine that -builds the LSE-on reduce until then. - -**What closes it:** retiring the C engine, or mirroring the rows above and -adding an `emit_lse=True` config per arch to the emitters. Done only when -`tools/check_byte_identity.py` is GREEN for every family at llvm20, llvm22 and -llvm23. - -**Sign-off:** avirgoel (owner), 2026-10-20. +| `emit_lse` on the unified problem, the `_reject_lse` support gate, and the gfx942 4-warp cohort exclusion (`library/kernels/common/attention_unified.py`: `_gfx942_4warp_eligible` is false under LSE, so the D256 bf16 and D128 sliding-window cohorts take the default 2D builder) | `rocke_unified_attn_problem_t`; the C selectors have no support gates yet, and `d256_gfx942_fast` (`attention_unified_selectors.cpp`) has no LSE clause. No C selector sets `emit_lse`, so a C-selected spec is always LSE-off and C output is unaffected | + +Out of scope here: the gfx950 tiled 2D C twin has no additive-bias params at +all (the Python builder's `use_additive_bias` has no field on the C spec), so +additive bias + LSE is built by Python only. That belongs to the additive-bias +gap (5). + +**Why deferred:** the C entry points take selector output, and no C caller +requests LSE today; the Python host and dispatcher are the only request path. + +**What closes it:** an `emit_lse` field on `rocke_unified_attn_problem_t`, the +support-gate refusals mirrored into the C selectors, `!emit_lse` mirrored into +`d256_gfx942_fast` (without it the C selector would keep the 4-warp cohort's +tile size for an LSE request while Python picks the default builder's spec), +and selector parity cases with LSE on. + +**Sign-off:** avirgoel (owner), 2026-10-20 (reduce, Python-only); 2026-10-28 +(reduce and tiled 2D kernels mirrored in C; request side remains). --- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/portable_ir/src/abi.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/portable_ir/src/abi.py index 69f677f3ab5d..66cec46c0f9f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/portable_ir/src/abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/portable_ir/src/abi.py @@ -75,7 +75,7 @@ #: is checked at library load: ctypes argtypes are a hand-written mirror of the C #: signatures, and calling a mismatched library with them is memory-unsafe rather #: than merely wrong. -BINARY_ABI = 5 +BINARY_ABI = 6 class AbiError(RuntimeError): diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention.json index e3653844ae60..09c1999b73d4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256/attention.json @@ -29,6 +29,30 @@ "sha256": "db218136b634c0011f18f4171fe5c1fe054c0a31f60fcb0e4bf606eba14b9fac", "status": "ok" }, + "attention/gfx942/2d_lse_mfma32_bf16_d128": { + "arch": "gfx942", + "bytes": 529371, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_lse", + "sha256": "76789686526efcedd42778e95ce53c8b8607447397d1e0306c3243079f35bc26", + "status": "ok" + }, + "attention/gfx942/2d_lse_striped16_bf16_d128": { + "arch": "gfx942", + "bytes": 399667, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "cdf465e6eaf51cc27605e081f72f49674235e9d56f98055b01ea3086d9b499be", + "status": "ok" + }, + "attention/gfx942/2d_lse_transposed32_bf16_d128": { + "arch": "gfx942", + "bytes": 488248, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_stqk_lse", + "sha256": "f4dfde43b327a95cdcc569ef1dac4f72a9cf978ebf23ad3641f64ee265dfb6a5", + "status": "ok" + }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", "bytes": 355976, @@ -71,10 +95,10 @@ }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10497, + "bytes": 10583, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "09854df6ad25883f6bf6c43eecb394c9c210b244924394f6976c308dcdad13b0", + "sha256": "eeb8ba17bd6b9f664e9a02113b53bb913d9b8fdb8d59179df0a64166ff79bd18", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -197,6 +221,30 @@ "sha256": "bb09632c2225d223d2daec4f586e9fa9a662502b11d1cb8bfb31a77b83a11e29", "status": "ok" }, + "attention/gfx950/2d_lse_mfma32_bf16_d128": { + "arch": "gfx950", + "bytes": 314684, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_lse", + "sha256": "87f12b31db45cc10dc85df48281501e10721f500259b6a55b5dda86ab6eac9eb", + "status": "ok" + }, + "attention/gfx950/2d_lse_striped16_bf16_d128": { + "arch": "gfx950", + "bytes": 171440, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "5aeedb828a6fa25fff9e19e90e54d24c3aa08d13922127805ec341bdefeaa62d", + "status": "ok" + }, + "attention/gfx950/2d_lse_transposed32_bf16_d128": { + "arch": "gfx950", + "bytes": 352854, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_stqk_s1_lse", + "sha256": "332dc60ad88293938e57dc4194c2483c21e12721603ac37be74d816eabf7e64e", + "status": "ok" + }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", "bytes": 161417, @@ -255,10 +303,10 @@ }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10462, + "bytes": 10548, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "aeb2bb788bd933c14b9ef4b6098ce2b14da748e9e501cbe3a87454b1f8a2cc35", + "sha256": "3e28886dc22e9f20989a4ecaedcc388a6a38a60d8822f5a7f39d303f5515d97f", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { @@ -362,6 +410,30 @@ "sha256": "2957a222d704be0e25ef149a256bec9203d4330358406ff8ef9610fa75d4ab03", "status": "ok" }, + "attention/gfx942/2d_lse_mfma32_bf16_d128": { + "arch": "gfx942", + "bytes": 529471, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_lse", + "sha256": "eacec37309a870aa88f00bd154b0fd7bd289918bf09fc83c9cebdb0ca71ef3ec", + "status": "ok" + }, + "attention/gfx942/2d_lse_striped16_bf16_d128": { + "arch": "gfx942", + "bytes": 399763, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "66c78d3b143707697544f2e966bcd4fe1001d535f400a695c9b18546edcfb69b", + "status": "ok" + }, + "attention/gfx942/2d_lse_transposed32_bf16_d128": { + "arch": "gfx942", + "bytes": 488348, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_stqk_lse", + "sha256": "28964d416187a795bc8ffa73136cee4beba4399c0671e0e84a8d949acd61348e", + "status": "ok" + }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", "bytes": 356076, @@ -404,10 +476,10 @@ }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10504, + "bytes": 10590, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "edbfaa7b91c851b499b0cc4fac7ceaa9e3c4922e60dd6ff464bc18f251f37843", + "sha256": "9f208dd3810d1e9712135bd2d022106e1bb7352290c18e8eb21d11bbc76fbb4f", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -530,6 +602,30 @@ "sha256": "caf014e9bad532cde9c3820ed188cdc6ff56df20ed3f8d1796a1a81fa14f210f", "status": "ok" }, + "attention/gfx950/2d_lse_mfma32_bf16_d128": { + "arch": "gfx950", + "bytes": 314780, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_lse", + "sha256": "1b02b801e28881de8dd40b0994d533573d780ebb91738de7ce07d02ed6e201f5", + "status": "ok" + }, + "attention/gfx950/2d_lse_striped16_bf16_d128": { + "arch": "gfx950", + "bytes": 171536, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "ed1a4d6bac001f05df0bbd02c75f734edd6a0eb58513a4ff811ce46362ee90c2", + "status": "ok" + }, + "attention/gfx950/2d_lse_transposed32_bf16_d128": { + "arch": "gfx950", + "bytes": 352950, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_stqk_s1_lse", + "sha256": "8083406611e0cce60ef09da21e842a7b12a46a20830e7b4ed30d2075d33e3753", + "status": "ok" + }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", "bytes": 161513, @@ -588,10 +684,10 @@ }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10469, + "bytes": 10555, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "ac383a27abbc3ce0195b48e843e12175f598ad1b1a4a5fe1f51641442c202750", + "sha256": "93085cd44fd95418414f6dc4cfbdce7d93c8ab6fb4961baba927c9bac5306d36", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { @@ -695,6 +791,30 @@ "sha256": "87fc31fbe4ecc0c55a463c3e1de856040a628ae723cb51515fd5748778709856", "status": "ok" }, + "attention/gfx942/2d_lse_mfma32_bf16_d128": { + "arch": "gfx942", + "bytes": 529535, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_lse", + "sha256": "5e9d09a0d5305df2237666634f17772352dd6df719ed0a5fa24fb1a3ffe9c1ce", + "status": "ok" + }, + "attention/gfx942/2d_lse_striped16_bf16_d128": { + "arch": "gfx942", + "bytes": 399827, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "f6cc646fe6d445b9af784cb1268d1cd6878d0724239cdb2427c30eed45b4f55e", + "status": "ok" + }, + "attention/gfx942/2d_lse_transposed32_bf16_d128": { + "arch": "gfx942", + "bytes": 488412, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32x8_stqk_lse", + "sha256": "eb9dd560dde0f9299fb58110dcf4c8dda71d376c03869f7f10c8045c7db75de9", + "status": "ok" + }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", "bytes": 356140, @@ -737,10 +857,10 @@ }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10568, + "bytes": 10654, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "38547bbde96c3f7fea0656216d3cda7814e8076218c51d756bb1d2f4000dc296", + "sha256": "b3bb8432873a068d142ab5d3eacca0dff88e0f97c9d4c88070aaa634ed6c8bf5", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -863,6 +983,30 @@ "sha256": "2f9446585481aa76c6a8074e5212a1408958d62ff44bf3344cdb5f165d7584e8", "status": "ok" }, + "attention/gfx950/2d_lse_mfma32_bf16_d128": { + "arch": "gfx950", + "bytes": 314844, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_lse", + "sha256": "d8d2009e3e81717e739b4ba32997f2b94d28983e5e0fb0d25343c75f576874c0", + "status": "ok" + }, + "attention/gfx950/2d_lse_striped16_bf16_d128": { + "arch": "gfx950", + "bytes": 171600, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_lse", + "sha256": "bf5f5cbd18c705680d81a4206941df460e6d9cddded09b2e0a3338d90690df73", + "status": "ok" + }, + "attention/gfx950/2d_lse_transposed32_bf16_d128": { + "arch": "gfx950", + "bytes": 353014, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16_mw32_mfma32_stqk_s1_lse", + "sha256": "d68a9bb5b0fc36915d914b18041e66c8fbab807a4c4734dc973a886007a83f69", + "status": "ok" + }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", "bytes": 161577, @@ -921,10 +1065,10 @@ }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10533, + "bytes": 10619, "family": "attention", "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "7518fc3a768b5f4598c9827648481455b523eb21bcf1e1f8b8c6134c30b195ae", + "sha256": "d7f993576194612f2c4beef0b5c0855e1a9502c8e2cc54a724386e69ee498963", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 9650c3fcab6d..029c281de181 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -1274,6 +1274,11 @@ "10": "17113cb8e55e17ac9bc3a37bc7144b91def22800f351f9074d66340dc50c73f8", "11": "f3475837408ae4290a05103719b938b01e04f7d48c7b401da58e5d21c2051172", "12": "6227b7f68633bb7fdec851a0d95337283504737a41b880ef1a8ee6d24527c6ae", + "13": "86b8b87bebefdb7322ebf7d0ff097a60a1ec1412d180a4a5cc0c5f8c2a32c07d", + "14": "93fb6779d998c4ea3383902186d2d852da356cceb62400ac7b09702af08554a1", + "15": "17eddc84585316dbf4c0f2dc8898e0d8eb7670e2c413ce06b10b36eab20e9873", + "16": "f1ca3ab0afb8378e3c569c588d41a8f8ff7fad1efdfaa3cae55380ee4994fdd4", + "17": "f20ab9afac7ee5291e189e835d368125d581984a1ecc41fe108507ed66c32848", "2": "57fcd1bb0c2dea494f172830d0ee0e266bfdb12770da8aef467f7582cf4b3015", "3": "7a1e7cf84ff9ed0a756b4cd03d192deebc21c4413a0abfd987de08e55ad543e9", "4": "399f51bf6a68669b35e0245b0c707a05aae3bc8e7814e163b2d439778ab7b9df", @@ -1350,6 +1355,9 @@ "18": "48d51c38964686d05421a90ba7ef9778fc1f3cb53ba200337ee66f761053bd04", "19": "520a827d9968e9fe8eea5be9898b4730e6bcfb60573818b86fbaa69cae7706c4", "2": "82cdca8f08ee9612cf5ea4f3ac729e826d283615ce0038eda2ce983c87bafd93", + "20": "faab3a2be33001ad480303f9ba04414e564ec4e98ebfc9f9a32bbc2b1d7733e1", + "21": "82c3c72998cd823e2637f88b0a1c72f348eb967f0614731fb61c69957d00c4dd", + "22": "619c8831f09e24ec84c6e11a4be906ba4b9a1ca485426056b0d5bebbd8b861d5", "3": "0bc5cf7cbe950277ca54549cb078916678e75b4db7a3880e22328478d89ffb3d", "4": "168b4be8598882c4beb38e84029c71293ab2f86538c4f8290c8c25b76a9634ef", "5": "a7494ff8169f6e0c7788dfd9321e6878b87226448c14d838d1fd7a7135d05b6c", @@ -1420,6 +1428,10 @@ "61": "6b9a1c2d6fd249dce2a15765149966ca576bebc58bffae91b315a8bd8bfc4e48", "62": "4545de951e607ad98a43162b8c4eeeab1a1a250a56164956641c27b13c7149f7", "63": "ae5a441ed96318b7bd047b5dccc1c5965b2a29c5752d917bc8d6f8415a6b098f", + "64": "dc5d14cafdf03825f97769fb812e0e405b978c77b99acc1d1b65c9e576ff42ee", + "65": "718e81bcd09bf25b0693a4afcc6dadb399b954bf0ce06e28aaf58648e62db238", + "66": "9ae3a8a167e75039158e8dc832ad76425a0863cd39c9c0ba2995c5c57604691d", + "67": "525634bba71dda9f5983ed8173b76ef2fd1f8e94cda7f9bafb013a627896dd17", "7": "72c9bcab5c89e84b8002b88b6ec2bea52850d0f9cbf824e56a9929621439b600", "8": "fcdd2a56823f8fd2f40d6f819ae0e82e4737eb43857db16a6a32a92fae5fde52", "9": "2fab4f3aea1a932eb3cd6e06aef8f9bc3296309f7b8a7bbf5adddea61df79b4e" @@ -1430,7 +1442,8 @@ "2": "d99a2e6b39ea9d9b463d9516ec55ec2ef8a0d4f1204c06517d4a8389682ddc20", "3": "1f4d7605f2cdcc4a94ca60aad0b10a45e66b677bb862554cba0d0dc45334d901", "4": "4b2f221c41bea34984f2156fb611158ea1c2fb4289e7ac4f44ca7b91ef696d1f", - "5": "1f480e98c92293827d4511c82d9b79baf45ab4df57bda8840a46fc64801ba28f" + "5": "1f480e98c92293827d4511c82d9b79baf45ab4df57bda8840a46fc64801ba28f", + "6": "093eea2640fa1b5b359284007f678732cae2da7b83fae90a772977b61616b8c4" }, "gfx950_attention_tiled_3d": { "0": "63c8cefe4601631d737cc00b2a88c787f723f2ba33cab0223da78d809c6c428c", @@ -1443,6 +1456,8 @@ "15": "77cbc7e611c9046832a959c50392299f8f007f65024a66d6fa6fd9d5ce3bef36", "16": "ac768048c589cb26f18c8ac4a0ac29ad00f109c591239fd601ca1803eaa458c6", "17": "6b5481c778882a5ee6110d0cac7e764c91d9f0028f95cf91864502ef91625d86", + "18": "2cee2330772fb5d3b34aee193719650d5419f1e47c46bb4a504eeec2587df6bb", + "19": "63c8cefe4601631d737cc00b2a88c787f723f2ba33cab0223da78d809c6c428c", "2": "7ca8ac489d8db2ad713d0c94d5e2b5814a38e1df54d82105e8ba1232ed6d1115", "3": "030fa0f8de882d8ecf675814aeed194d744c6642345c2e1f87a736f386a5ff7b", "4": "206a5d7824acf278b093b00d7b26723a9b2e6815b66e097d64296c4df3518ec9", @@ -3029,6 +3044,11 @@ "10": "8a6ef4f23439f6903da0569076032ababcf2a98ab98e5f2aa5abc2ddefec3602", "11": "7cdca35b8857df8ce9ab70ef1532ed171438dd1bc5dd8edbc88db35fb652e80d", "12": "f45f3e00722910bed025b8357ea96cb9da4cd9ea5d453ad14cd1669f332a8a3e", + "13": "f40014e952c10d7e1a2d0d8716c8633d2035800642361e5d462112016b6cdccb", + "14": "704c01b76f73b57f4a0305a65b2ae4f7f1058b7c8f12c8ce1ae405ec94b7f286", + "15": "0c9c2dc50bdfab35ac52d045e126c87105ba32cac0e848ccf560ee0b2e78b53e", + "16": "c1a89a4542a1a394f44322adb4d5a8bae37bb9de636770025f3092a39a6f47e0", + "17": "b020620e6517037925ad70a8ca71aa7548f785b97bab3601fa169dfc8d410894", "2": "d41737466147aa059046db9536fb476d681f996a5028a7b465a8b7ceb14177ee", "3": "1914a01617f907e4790693583ea5155d34af6ca09ca0ba62135d70fd31209a86", "4": "8ce860e6181878106830c77f55b195b57a1ec011696a3a2795f0498a6004edfa", @@ -3105,6 +3125,9 @@ "18": "8ea2760245c71f79c9b42ae1d10362ef2a07df38ba65e17de924d813592aef1b", "19": "ae54f0e8c4ab82f253693ccf4077a4a30bcdd8bf4b0fe965a94357eb3b7136d4", "2": "d776e66b764f75793fc25b17e2639b306c289592fff73953fcede2e76b189322", + "20": "7cd9da7ec31fff72f25b3c5fd8195535771ee48ed073fb72e89895e2495a3a99", + "21": "806b56ba66837f09e3f56f4ff9ec3603a6f9ecbca388d0da00fb5fd573320247", + "22": "197c0bbd20666e3b52b0edd1c5d2d0c08261b7aeec06c5c05c565fdbad24d9ac", "3": "41f971cb9da04b95343d06bd7c15143ff7cf3e82eb9baee619b563bc672d3838", "4": "5eff7adc1f3424caab3164b362e805f6561e65955071ba010deb46f6a8aa33a2", "5": "26bb8e65bf6a98aa64a294fc00756b12910f3f2a2c4a3c7e12a7a05d5f0c0244", @@ -3175,6 +3198,10 @@ "61": "c4b7fdfadfd92354463ad1fa13a4dd17d87b32acd92ba14b47be644cbc9412ce", "62": "32d908604514f4fa2c094adfd67ab0313fedc953c0b3feba7ea19bcb2ae87dd7", "63": "0a9d04484c6a56db9a6e25c55ec4998a788dd4904b3f5791cd6b58d99eec9dd1", + "64": "da27a564535e6dd9991ae53a40410a5d83f47fe08cce29dfdcf49547109c9c49", + "65": "b6d959cef65dbb71e1c5ce48848a5da8c028917db52a140bcc696d9f69d3317c", + "66": "af2edf8fdbf7280defcfe5b2a9afb508299cbf91bf1e6d99c06757c00805c880", + "67": "8d29bb3ea04a7ff8d41404cff67668feaea97414785c186b5c1b553dcf08d4d9", "7": "de980c05bb12883d8e1b5e1e2d3233f38a21cb7cb41faf64881c637139f6aef4", "8": "da07a09a050011bb4663845c27db729ba96875ee4ba89bc74191cf974b6db1bb", "9": "112f48a27fbfc7513d6fef46d1e6cae4f0212f2f4bee0617104307fda5a9b836" @@ -3185,7 +3212,8 @@ "2": "fb726c7366cf45b6c6893fb5fa4610520f7aa61864ee6f2b69e44776588dd2fa", "3": "50b3aca4dbd9b425ecb92bc9cf908be84e3ebc91dcb32879fa7c205d196c6681", "4": "0bcce2d8ca0a9d51aab0a1e776dc7de5b6c2fe662180642a46cab54c5f5f255c", - "5": "7045a67c8fb900d078d3fa3ae38907661912ccacaf6b080a16958074eaea21e0" + "5": "7045a67c8fb900d078d3fa3ae38907661912ccacaf6b080a16958074eaea21e0", + "6": "570c4873db463d1152be29fe809068c5bfc0d7f8c3d574a28370a196614d3cdf" }, "gfx950_attention_tiled_3d": { "0": "211ac134586448e278d5662ad89a605338b755e91bf63d1584e267dfc79254bf", @@ -3198,6 +3226,8 @@ "15": "5922ae9e30f10d84ba06293cf1b9f530872f26ac4b31d264ee0e7e41d96f85f2", "16": "11c86b374923581855b8fbad8e2e0728c03c86085f9ad5e01f88bb0dc18cc510", "17": "53d870b45bf1beb8bc838b6941a732fc425feccf36347b01a59ec93cca6c06c4", + "18": "89b0866d30ee69cf00927ffe5cf6f3117d8e99269185ed21de78d508d284b7f0", + "19": "360bd847b011636cbd63eb70c90cda173cd2c7c913a16abc881f53dfbbaacfc0", "2": "43225f4bde00693107c9eba8bb29ede7589c65b60cee13b6e310bce7260bdce2", "3": "011317d12e0cd324dfdf21438343e02c900930fb3c5ccd02ec2752f73188734b", "4": "68468fbf2b4c26089c675cc417e8af43b308a6a7838e9604853845fa154e6dd4", @@ -4783,6 +4813,11 @@ "10": "258774f7a4f3430c5ab90760035918f0a6759e0f9fc0c62f28e32626de3a85aa", "11": "da498a36c8a45e8d822f67d528fdbb528915a3ace93dc148f52f8b48141ecf90", "12": "c11599d44791da28410a5af9bf6eed88d3a77bb88e602969ca7aa8eeb3f9d2cc", + "13": "b550bff27d546914222eda35bef90306ae128d60f01db4b6290ce61524451747", + "14": "a7c78d9eb2df50c53c6246b1b8baa1379e6a39cdfa4d8e2d221fa91f035314bf", + "15": "17494678dbd9a4f25126df82c23994f0b001367f056a730141069fbc3e8ba0ee", + "16": "53968737258753d6f84ba1b9e95be559f894f84c44376a9277221e74a23c86a8", + "17": "23550b712bce1643e488b65478b006cbb3d3c8e5e96a573b2e5ef5447a385273", "2": "39f5bfca25beafff8a1634549fa92f2512fc6baa207ee2891db0789c540b7a22", "3": "09bf31f3d9c5883bee0c44bf01b49851e43ad800fe3477eb30bd75b5928ed0d2", "4": "fda6cb477cc284b634b74a5a3553e38517964d40d2aa3169d0d3a755e89426fa", @@ -4859,6 +4894,9 @@ "18": "6bbf65ce27b55559b83aa0de86b36c691b3792997081f047a48e00b575c974ae", "19": "7b0041c22c25865cfb15ffe70640780a5fa79f5f6fe44781bd9858579d36b7eb", "2": "56153f0bef9d8262404f7d0887de981b0bd919840bc2310b6473b98314df26aa", + "20": "a57d936897c12305b53746e75e566c805d8fbbce9d5eba992df17c89bb759fbc", + "21": "a29cfc841cbe873fd324a01cf60e0c35d255331577fdafaf40942bac9c35e05a", + "22": "b272e8028eadf023648b3cfbaec0119de8aa4621326a02ccdf2956c691e29997", "3": "9398d2999e4caa2d1fa3ed7be10683b884fb813f22eb0404bb9cf8d472575904", "4": "e5eddf3ecfd8fae5d62b08af12677b6d15fd9a1111225b4a08ea828cb12d0905", "5": "523279f91929d6188a955c2490f810830fc15abf00df25dc055c093af8db6da5", @@ -4929,6 +4967,10 @@ "61": "1c2c15cc564ea3b516ef844207a7ca6f544fe153e8287c8619b2d6c8d6b04e41", "62": "d1712525752acf3543932485fde5cc6e5f630128030aeee43f1ac94b20c03e95", "63": "a99d1b8c3587031b9848350bbbd3cb8e2e75ed5efbaecbda0fc4a3a15b6be128", + "64": "3ce7d597a2e9ecd6f1305a68f0c46e354cb12385e23a0a39e9c9f482481f76fd", + "65": "fcd7ccce7f6e9d174a78dafe3382e668b488c7be715653ac3308dd963176d3fc", + "66": "ded6524f2739af4192f3bf67ea5dc4ff059d39c36124ed6f01c6f855b92ef716", + "67": "dd449e63aa2905635f3fd7352cfd467acfeaf8fc93b244665f3a04f48f293b83", "7": "29103ea75c671cc7a30350a778f16146c88f7f05a8aca8300afae7fa0dbcbe99", "8": "9f46cdbb82b854e5d3caa8d19749469bfd1884f12bf96f8ec0e95bbe6b16ebbd", "9": "9d3ca2669adc1f687172bc70053dd47c37e11606af25be19869b2ff3203ecbcd" @@ -4939,7 +4981,8 @@ "2": "251324d112d6c528566f26147e83d1d17500fb7b2d960c4f63e8a686501cd0bc", "3": "38699dcda4a34bad75630e11de51cd1fd9105a8bc91967cd3e41897292bc6a10", "4": "999975e55cbecdfcc0160fdb0eb2ae74b3ea84a36c30673580a464ea1ee2cf86", - "5": "8a030aa0a4d71c252cfc14adb5984ac2dbecd77c3bfcf1da1b9474db26cffc1c" + "5": "8a030aa0a4d71c252cfc14adb5984ac2dbecd77c3bfcf1da1b9474db26cffc1c", + "6": "8c9b365749845364c858982a743c0b78ff94260b6a8d9172b3dc3900ed15b17f" }, "gfx950_attention_tiled_3d": { "0": "22c3c50dcdeecfdd1706c2c3806786e493f4576336bd5960ed9c9029bb8ee919", @@ -4952,6 +4995,8 @@ "15": "af9726d24b47db9c198912f2d3f6ecfd2701cb3a720da6f3619d0949abd9f6bf", "16": "6ed051c1af811770c2af798cdca13056c07fce94afc67d697e71f69796356fcc", "17": "813cae2f8db1cde8de40cde864c8f67fb51e2e5b3cf3b5d0b128b306934c4695", + "18": "bb08fbbeb808eba61d61dc2b0c0be758a7ef326d5ffaf6cd25cc42c9d9a2381d", + "19": "ccbb69e46dbefc3b1b36c7d2a838206534a4afe856c29adfa96ff31f8fe7e558", "2": "f0a777ee162213e55e259419d46d1791a975efb8fb0bf272d7536cd9d9efaeca", "3": "f92162681e4e1b1c3bd8ff591baefe4dcfa3a840ade95e46c4b3ae47d3498880", "4": "41bf0207a6966ba320af8c251bfc51a0791e82f3782b204b3cc1e07e2c265a57", @@ -6664,6 +6709,11 @@ "10": "1925a2ff801c2d367abbe1ecece08384dd38298279ca41b2a1af34bf9e06b25c", "11": "cfede622a546315810cee66a584e7a62cf05df8f7b0da560c707a7e65e5926a2", "12": "8327bf97823a7823a95536e84d14ca722bf5b8c32f7118d5f310d5ae1eeb98a9", + "13": "a0720903a7d354e085d4c68189279ee6eacfff930b865499707d59d91fe178bd", + "14": "487415ff6b42e500ee274f222930de09569f4d396996d9dcc1aa168e38798797", + "15": "d3a693c511a915d5706a56d83611d666a8b5a97746f59d124fd3345c012327b2", + "16": "8c613582fb8777cbf3a25597681dec2cc5d94a110bcceb3e6d3ea538bb58ffd3", + "17": "bf20d61d110b716e4c42c6e3c73cc329e75261ed9e63bc7f76b9f3c42efd4613", "2": "8c4bcf8012621b5ea67bf058152b4a680dcc66a1641babdb4f6529b13bea9136", "3": "73acf9cd47d2e5e4d456413abc81d2d8717f89a50236a6a73a438d95f5f46352", "4": "638b0330c24b321de32314cce43e0a2dad8ef1c387ea74eda7a65de4573529cc", @@ -6740,6 +6790,9 @@ "18": "8d593285bfb858ad7385ab547d1d809c626986a8aafdc2d54d88d30ce4a3863d", "19": "65858a9f795e872b15652d116a9f1ebea6d8e3b9e8c981910785ed1913fe5256", "2": "6f4de2c5041fdca3d24abc5d74cb661146e716ad06b113600cefd059875f26c1", + "20": "13475d55ceeaa04216a6a73d39395ba262ce1fe0a285fda57bb4c969e1ffc93a", + "21": "8f5e4de1a6939da30edda46461e70e3c8293e45ae6bcd85974a8b4b93962087b", + "22": "3c69e5d2b5f6c817885851dccd04ecc66192a8264720c1ea0e0063407b02a994", "3": "e8153af8740777bab3d5613e33d52b40b4616cf4b1890a0559b76c0eb31e2d85", "4": "3e17cdf6ffb2daa7a7c457c1741f3a6350a5939897d5d7b6e9bc21fa64ca60e0", "5": "dff48e7a6319aae6c1f0a67bc8eea2e79686a4caaf34b9cb1704cc62d2574f67", @@ -6810,6 +6863,10 @@ "61": "bf31c9c41afd790d5400e98f31dbc19da1d5f22dd1cb62d6ecc3bdf4ca480b9d", "62": "1b5ec55411b01a7a7e00ebfedc73dfa66ad01ffc2501410f671893aea4c9e6c1", "63": "de1c125baa190df0aa53f14e6d730b8728cd33dfc0cb32387ba392b88c99871c", + "64": "f4c6ce377965f2fb53d583e7ce80923b9dfd462e325fd9e6c92d24f2a03d110e", + "65": "378c535e5425529c682f6c5bb87675241a425bc7e4334b7a50a94e5d673be96a", + "66": "7383690c0f04c4302ae9f36c106473b4bec89e99ba2d60e65a1b85d1fd8e5c10", + "67": "2ee389bac285bba8618de8afa3947939c246ef70f37315ea6894a28bf6a4a9a7", "7": "71885a9b807d71cec0b1a35eee2efeed8b7f6ccb9bc503ba701a4f957ecedc6e", "8": "dfda5480270460c3cff1a887485ca9517356a2c34bcb4050aad8d789ed62c293", "9": "d1c32503e3084dde749f4b40f6468cc4227f577debca994033e07bbbe54d5f1a" @@ -6820,7 +6877,8 @@ "2": "361a82466d9c33eb0c50d5d511c6e0c831ae3c317f192752c31bea4bd715829b", "3": "1af7b32a44343599433ae0a320f28437014e4f00344578a4602d97e36a4a0466", "4": "b7deee184c7b806981600cc7300b89d9653426bb3a85dd220dde96af788fb874", - "5": "7e8f4056da548709d5ab9d17edee524ba06ff2e16497198f9f0aae8a76735e2c" + "5": "7e8f4056da548709d5ab9d17edee524ba06ff2e16497198f9f0aae8a76735e2c", + "6": "4c890c3b268497e172919e8be183533f933d4602cbe9f0ca4f8b46aee0736a4f" }, "gfx950_attention_tiled_3d": { "0": "2d15cfdb485bc112c4364d6b000586b3da7cfa0403b170a7f4c4662a3003ab84", @@ -6833,6 +6891,8 @@ "15": "942c3c72cc504c539d4abc6201db2bab548871427f35a335258eee0614c1491e", "16": "5276267bc23de9579470828f4d0452b249d301e70727437f315544ae9e3ed9c0", "17": "99acbf37277a473e4cedd12d0595fc6b304f7af70c39fae62f08c0eb1a0a13a6", + "18": "52f9076af863cd1f6062f112db555f90d28a58df1a610dde09067b2fdc52f66a", + "19": "6d33f51f39c8d4a54c1be3ca2098dea8cce826d7395a0ecb8a09417e80d98f64", "2": "1680787ea992fe51bc10f0d663dce835be3c0dc742b4a36023964a9c1e671ebf", "3": "d543f94c41159c5af6822c0656497b16bcada09f475aa2205c2ed937becaa526", "4": "e3c11f0bea0406fd826b858b25c0273db7685e0fc1d04cc8f90c4789fa474a9f", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 587e8d829511..c8be1742c44f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -3537,6 +3537,37 @@ def add(family, case_id, arch, build): emit_lse=True, ), ) + # Tiled 2D with LSE output, one case per epilogue branch (each branch has + # its own LSE store site). + for _branch, _branch_kw in ( + ( + "transposed32", + dict( + use_mfma_32x32=True, + block_m_per_warp=32, + use_transposed_qk_32x32=True, + use_transposed_scalar_state=True, + ), + ), + ("mfma32", dict(use_mfma_32x32=True, block_m_per_warp=32)), + ("striped16", {}), + ): + add( + "attention", + f"attention/gfx950/2d_lse_{_branch}_bf16_d128", + "gfx950", + build_attention_2d( + f"irhash_attn_950_2d_lse_{_branch}_bf16", + "gfx950", + head_size=128, + block_size=64, + num_query_heads=4, + num_kv_heads=2, + dtype="bf16", + emit_lse=True, + **_branch_kw, + ), + ) add( "attention", "attention/gfx942/2d_fp16_d128_b64", @@ -3586,6 +3617,37 @@ def add(family, case_id, arch, build): use_i64_kv_addr=True, ), ) + # LSE-on tiled 2D, one case per epilogue branch (transposed 32x32x8 direct + # store, non-transposed 32x32x8 via Acc_lds, narrow 16x16 striped). The + # LSE-off cases above stay byte-identical. + for _branch, _branch_kw in ( + ( + "transposed32", + dict( + use_mfma_32x32x8=True, + use_transposed_qk_32x32=True, + block_m_per_warp=32, + ), + ), + ("mfma32", dict(use_mfma_32x32x8=True, block_m_per_warp=32)), + ("striped16", {}), + ): + add( + "attention", + f"attention/gfx942/2d_lse_{_branch}_bf16_d128", + "gfx942", + build_attention_2d( + f"irhash_attn_942_2d_lse_{_branch}_bf16", + "gfx942", + head_size=128, + block_size=64, + num_query_heads=4, + num_kv_heads=2, + dtype="bf16", + emit_lse=True, + **_branch_kw, + ), + ) add( "attention", "attention/gfx942/3d_fp16_d128_b64",