From a1263fadf3c3149752b8d286ff35464bd445ff08 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:28:34 +0800 Subject: [PATCH] fused_qkv_rope in UOp try 2 (#17619) * fused_qkv_rope in UOp try 2 * dont need that * less --- extra/thunder/amd/fa.py | 37 +++++++++++---- extra/thunder/amd/fused_qkv_rope.cpp | 69 ---------------------------- test/backend/test_asm_gemm.py | 2 + test/backend/test_llama_kernels.py | 19 ++++---- 4 files changed, 38 insertions(+), 89 deletions(-) delete mode 100644 extra/thunder/amd/fused_qkv_rope.cpp diff --git a/extra/thunder/amd/fa.py b/extra/thunder/amd/fa.py index 685a1b8e77..36cdfad824 100644 --- a/extra/thunder/amd/fa.py +++ b/extra/thunder/amd/fa.py @@ -19,16 +19,33 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None @functools.cache def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int): - code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text() - threads = 256 - thread_idx = UOp.special(threads, "lidx0") - block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1") - sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y, - arg=KernelInfo(name="fused_qkv_rope_forward")) - compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", - f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"] - lib = HIPCCCompiler(arch, compile_args).compile_cached(code) - return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib))) + group_size = H // H_KV + q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D) + xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D) + b, n = UOp.range(B, 0), UOp.range(N, 1) + pair = UOp.range(D // 2, 2) + even = pair * 2 + c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float) + s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float) + ordered:UOp|None = None + for kvh in range(H_KV): + q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v)) + x_in = xqkv.after(ordered) if ordered is not None else xqkv + stores:list[UOp] = [] + for rep in range(group_size): + a = x_in[b, n, kvh, rep, even].cast(dtypes.float) + bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float) + h = kvh * group_size + rep + stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))] + a = x_in[b, n, kvh, group_size, even].cast(dtypes.float) + bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float) + stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)), + k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)), + v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]), + v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])] + ordered = UOp.group(*stores) + assert ordered is not None + return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward")) @functools.cache def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp, diff --git a/extra/thunder/amd/fused_qkv_rope.cpp b/extra/thunder/amd/fused_qkv_rope.cpp deleted file mode 100644 index fcb7ee6bf9..0000000000 --- a/extra/thunder/amd/fused_qkv_rope.cpp +++ /dev/null @@ -1,69 +0,0 @@ -#include -#include - -#ifndef ATTN_B -#define ATTN_B 2 -#endif -#ifndef ATTN_N -#define ATTN_N 8192 -#endif -#ifndef ATTN_H -#define ATTN_H 32 -#endif -#ifndef ATTN_H_KV -#define ATTN_H_KV 8 -#endif -#ifndef ATTN_D -#define ATTN_D 128 -#endif -#ifndef THREADS_PER_BLOCK -#define THREADS_PER_BLOCK 256 -#endif - -constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; -constexpr int HALF_D = ATTN_D / 2; -constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D; - -extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void -fused_qkv_rope_forward( - __hip_bfloat16* __restrict__ q, - __hip_bfloat16* __restrict__ k, - __hip_bfloat16* __restrict__ v, - const __hip_bfloat16* __restrict__ xqkv, - const __hip_bfloat16* __restrict__ freqs_cis) { - const int b = blockIdx.x; - const int n = blockIdx.y; - const int bn = b * ATTN_N + n; - const int packed_bn = bn * ATTN_H_KV * PACKED_D; - const int q_bn = bn * ATTN_H * ATTN_D; - const int kv_bn = bn * ATTN_H_KV * ATTN_D; - - if (threadIdx.x < HALF_D) { - const int pair = threadIdx.x; - const int even = pair << 1; - const float c = static_cast(freqs_cis[((n * HALF_D + pair) * 2) + 0]); - const float s = static_cast(freqs_cis[((n * HALF_D + pair) * 2) + 1]); - - for (int kvh = 0; kvh < ATTN_H_KV; kvh++) { - const int base = packed_bn + kvh * PACKED_D; - - for (int rep = 0; rep < GROUP_SIZE; rep++) { - const int qbase = base + rep * ATTN_D; - const int h = kvh * GROUP_SIZE + rep; - const float a = static_cast(xqkv[qbase + even]); - const float bb = static_cast(xqkv[qbase + even + 1]); - const int out = q_bn + h * ATTN_D + even; - q[out] = static_cast<__hip_bfloat16>(a * c - bb * s); - q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c); - } - - const float a = static_cast(xqkv[base + GROUP_SIZE * ATTN_D + even]); - const float bb = static_cast(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]); - const int out = kv_bn + kvh * ATTN_D + even; - k[out] = static_cast<__hip_bfloat16>(a * c - bb * s); - k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c); - v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even]; - v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1]; - } - } -} diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index e18458f808..6e34b996c0 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -1,4 +1,5 @@ import unittest +import functools from tinygrad import Tensor, Device, dtypes, Context from tinygrad.helpers import getenv, system, DEV from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm @@ -9,6 +10,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX # Use DEV=NULL:HIP:gfx950 to also test the assembly def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") +@functools.cache def has_hipcc(): try: system("hipcc --version") except Exception: return False diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py index f03a57b393..ee8538c321 100644 --- a/test/backend/test_llama_kernels.py +++ b/test/backend/test_llama_kernels.py @@ -99,22 +99,20 @@ class TestLocalAmax(unittest.TestCase): assert_kernel_count(2) self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]]) -@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run") class TestFusedQKVRoPE(unittest.TestCase): SHAPE = (2, 8192, 32, 8, 128) + def setUp(self): + if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs") + def rand_bf16(self, *shape:int) -> Tensor: return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize() - def freqs_cis(self) -> Tensor: - _, N, _, _, D = self.SHAPE - return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize() - - def test_llama31_8b_forward(self): + def test_forward(self): Tensor.manual_seed(0) - B, N, H, H_KV, D = self.SHAPE + B, N, H, H_KV, D = 1, 32, 8, 2, 16 GROUP = H // H_KV - freqs_cis = self.freqs_cis() + freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize() x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D) q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D) @@ -131,12 +129,13 @@ class TestFusedQKVRoPE(unittest.TestCase): self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch") self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch") - def test_llama31_8b_backward(self): + @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile") + def test_llama31_8b(self): Tensor.manual_seed(1) B, N, H, H_KV, D = self.SHAPE PARTIALS = 2 GROUP = H // H_KV - freqs_cis = self.freqs_cis() + freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize() dq = self.rand_bf16(B, N, H, D) dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D) dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)