geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

fused_qkv_rope in UOp try 2 (#17619)

Browse Source 
* fused_qkv_rope in UOp try 2

* dont need that

* less
This commit is contained in:
qazal
2026-08-20 12:28:34 +09:00
committed by GitHub
parent e6324d1e1c
commit a1263fadf3
4 changed files with 38 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/thunder/amd/fa.py
+27 -10
View File
@@ -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,
extra/thunder/amd/fused_qkv_rope.cpp
-69
View File
@@ -1,69 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#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<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
const float s = static_cast<float>(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<float>(xqkv[qbase + even]);
const float bb = static_cast<float>(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<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
const float bb = static_cast<float>(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];
}
}
}
test/backend/test_asm_gemm.py
+2
View File
@@ -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
test/backend/test_llama_kernels.py
+9 -10
View File
@@ -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)