| """Experimental Metal 4 online-softmax attention with configurable Q/K tiles. |
| |
| Dense, noncausal, batch-one, equal-head attention only. No approximated sparsity. |
| Not selected by the runtime unless full-pipeline validation supports a speed win. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import error, measure |
| from h3_mlx.ops import attention |
| from h3_mlx.runtime import configure_memory, environment |
|
|
|
|
| @lru_cache(maxsize=None) |
| def build(bq, bk, groups): |
| return mx.fast.metal_kernel( |
| name=f"h3_flash_{bq}_{bk}_{groups}", |
| input_names=["q", "k", "v"], |
| output_names=["out"], |
| ensure_row_contiguous=False, |
| header=""" |
| #include <metal_tensor> |
| #include <MetalPerformancePrimitives/MetalPerformancePrimitives.h> |
| using namespace mpp::tensor_ops; |
| """, |
| source=r""" |
| const int N = q_shape[0], H = q_shape[1], D = 128; |
| int sg = thread_position_in_threadgroup.x / 32; |
| int row = (threadgroup_position_in_grid.x * SG + sg) * BQ; |
| int head = threadgroup_position_in_grid.y; |
| // SIMD groups own independent query tiles; the last group may be empty. |
| auto Q = tensor(const_cast<device T*>(q) + head*q_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(q_strides[0])}); |
| auto K = tensor(const_cast<device T*>(k) + head*k_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(k_strides[0])}); |
| auto V = tensor(const_cast<device T*>(v) + head*v_strides[1], dextents<int,2>(D,N), array<int,2>{1,int(v_strides[0])}); |
| auto O = tensor(out + head*D, dextents<int,2>(D,N), array<int,2>{1,H*D}); |
| if (row >= N) return; |
| constexpr auto qkdesc = matmul2d_descriptor(BQ,BK,128,false,true,false); |
| constexpr auto pvdesc = matmul2d_descriptor(BQ,128,BK,false,false,false,matmul2d_descriptor::mode::multiply_accumulate); |
| matmul2d<qkdesc,execution_simdgroup> qkop; |
| matmul2d<pvdesc,execution_simdgroup> pvop; |
| auto qtile = Q.template slice<128,dynamic_extent>(0,row); |
| auto ktile = K.template slice<128,dynamic_extent>(0,0); |
| auto vt = V.template slice<128,dynamic_extent>(0,0); |
| auto vtile = V.template slice<128,dynamic_extent>(0,0); |
| auto qs = qkop.template get_left_input_cooperative_tensor<T,T,float>(); |
| qs.load(qtile); |
| auto scores = qkop.template get_destination_cooperative_tensor<decltype(qtile),decltype(ktile),float>(); |
| auto maximum = qkop.template get_row_reduction_destination_cooperative_tensor<decltype(qtile),decltype(ktile),float>(); |
| auto denom = maximum; |
| auto tile_max = maximum; |
| auto tile_sum = maximum; |
| auto alpha = maximum; |
| auto prob = pvop.template get_left_input_cooperative_tensor<T,T,float>(); |
| auto p_shape = tensor(const_cast<device T*>(q), extents<int,BK,BQ>(), array<int,2>{1,BK}); |
| auto accum = pvop.template get_destination_cooperative_tensor<decltype(p_shape),decltype(vtile),float>(); |
| #pragma unroll |
| for (int i=0;i<maximum.get_capacity();++i) { maximum[i]=-INFINITY; denom[i]=0.0f; } |
| #pragma unroll |
| for (int i=0;i<accum.get_capacity();++i) accum[i]=0.0f; |
| threadgroup T scratch[SG*BQ*BK]; |
| auto P = tensor(scratch+sg*BQ*BK, extents<int,BK,BQ>(), array<int,2>{1,BK}); |
| for (int start=0;start<N;start+=BK) { |
| auto kk = K.template slice<128,dynamic_extent>(0,start); |
| qkop.run(qs,kk,scores); |
| #pragma unroll |
| for (auto it=scores.begin();it!=scores.end();++it) { |
| auto coord = scores.get_multidimensional_index(it); |
| *it = start+coord[0]<N ? *it*0.08838834764831845f : -INFINITY; |
| } |
| reduce_rows(scores,tile_max,reduction_operation::max,-INFINITY); |
| #pragma unroll |
| for (int i=0;i<maximum.get_capacity();++i) { |
| float next=max(maximum[i],tile_max[i]); |
| alpha[i]=exp(maximum[i]-next); maximum[i]=next; |
| } |
| #pragma unroll |
| for (auto it=scores.begin();it!=scores.end();++it) { |
| *it=exp(*it-*maximum.map_iterator(it)); |
| } |
| reduce_rows(scores,tile_sum,reduction_operation::sum,0.0f); |
| #pragma unroll |
| for (int i=0;i<denom.get_capacity();++i) denom[i]=denom[i]*alpha[i]+tile_sum[i]; |
| #pragma unroll |
| for (auto it=accum.begin();it!=accum.end();++it) *it *= *alpha.map_iterator(it); |
| // FP32 and BF16 cooperative destinations may have different lane |
| // layouts. Convert by logical coordinates, never by array index. |
| #pragma unroll |
| for (auto it=scores.begin();it!=scores.end();++it) { |
| if (scores.is_valid_element(it)) { |
| auto coord=scores.get_multidimensional_index(it); |
| scratch[sg*BQ*BK+coord[1]*BK+coord[0]]=T(*it); |
| } |
| } |
| simdgroup_barrier(mem_flags::mem_threadgroup); |
| prob.load(P); |
| auto vv = V.template slice<128,dynamic_extent>(0,start); |
| pvop.run(prob,vv,accum); |
| } |
| #pragma unroll |
| for (auto it=accum.begin();it!=accum.end();++it) *it /= *denom.map_iterator(it); |
| #pragma unroll |
| for (auto it=accum.begin();it!=accum.end();++it) { |
| if (accum.is_valid_element(it)) { |
| auto coord=accum.get_multidimensional_index(it); |
| if (row+coord[1]<N) out[((row+coord[1])*H+head)*D+coord[0]]=T(*it); |
| } |
| } |
| """, |
| ) |
|
|
|
|
| def flash(q, k, v, bq=16, bk=64, groups=4): |
| if ( |
| q.ndim != 4 |
| or q.shape[0] != 1 |
| or q.shape[-1] != 128 |
| or q.shape != k.shape |
| or q.shape != v.shape |
| ): |
| raise ValueError("Experimental attention requires equal [1,N,H,128] tensors") |
| n, h, d = q.shape[1:] |
| return build(bq, bk, groups)( |
| inputs=[q[0], k[0], v[0]], |
| template=[("T", q.dtype), ("BQ", bq), ("BK", bk), ("SG", groups)], |
| grid=(((n + bq * groups - 1) // (bq * groups)) * groups * 32, h, 1), |
| threadgroup=(groups * 32, 1, 1), |
| output_shapes=[(1, n, h, d)], |
| output_dtypes=[q.dtype], |
| )[0] |
|
|
|
|
| def main(a): |
| configure_memory() |
| mx.random.seed(731) |
| out = Path(a.output) |
| out.parent.mkdir(parents=True, exist_ok=True) |
| report = { |
| "environment": environment(), |
| "repeats": a.repeats, |
| "warmup": 2, |
| "scope": "Synthetic BF16 Q/K/V; dense online softmax. Different reduction/rounding may affect generated images.", |
| "trials": [], |
| } |
| for n in a.tokens: |
| q, k, v = ( |
| mx.random.normal((1, n, 56, 128)).astype(mx.bfloat16) for _ in range(3) |
| ) |
| mx.eval(q, k, v) |
| ref = attention(q, k, v) |
| mx.eval(ref) |
| variants = [(0, 0, 0)] + [ |
| (bq, bk, sg) for bq in (16, 32) for bk in (32, 64, 128) for sg in (2, 4) |
| ] |
| for bq, bk, sg in variants: |
| row = {"tokens": n, "bq": bq, "bk": bk, "simdgroups": sg} |
| try: |
| fn = ( |
| (lambda q=q, k=k, v=v: attention(q, k, v)) |
| if bq == 0 |
| else ( |
| lambda q=q, k=k, v=v, bq=bq, bk=bk, sg=sg: flash( |
| q, k, v, bq, bk, sg |
| ) |
| ) |
| ) |
| row.update(measure(fn, repeats=a.repeats), **error(ref, fn())) |
| except Exception as exc: |
| row["error"] = f"{type(exc).__name__}: {exc}" |
| report["trials"].append(row) |
| out.write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
| if "error" in row: |
| return |
| del q, k, v, ref |
| gc.collect() |
| mx.clear_cache() |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--tokens", nargs="+", type=int, default=[4181, 16469]) |
| p.add_argument("--repeats", type=int, default=5) |
| p.add_argument("--output", default="reports/optimization-v2/attention.json") |
| main(p.parse_args()) |
|
|
