mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 16:36:07 +00:00
llm: trim unused AMD quant paths
This commit is contained in:
@@ -1,66 +0,0 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
|
||||
from tinygrad import Device, Tensor, TinyJit
|
||||
from tinygrad.llm.kernels.amd import amd_flash_attention_decode, flash_attention_causal_cached
|
||||
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required")
|
||||
class TestAMDFlashAttention(unittest.TestCase):
|
||||
def _test_decode(self, max_kv_len:int, valid_kv_len:int, n_heads:int=16, n_kv_heads:int=2, quantized:bool=False):
|
||||
rng = np.random.default_rng(1)
|
||||
q_np = rng.standard_normal((1, n_heads, 1, 256)).astype(np.float16)
|
||||
kv_np = rng.standard_normal((2, 1, n_kv_heads, max_kv_len, 256)).astype(np.float16)
|
||||
scale_np = np.maximum(np.max(np.abs(kv_np.astype(np.float32)), axis=-1), 1e-8) / 127
|
||||
if quantized:
|
||||
kv_np = np.clip(np.rint(kv_np.astype(np.float32) / scale_np[..., None]), -127, 127).astype(np.int8)
|
||||
q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize()
|
||||
scale = Tensor(scale_np.astype(np.float16)).realize() if quantized else None
|
||||
|
||||
@TinyJit
|
||||
def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len, scale).realize()
|
||||
|
||||
out = None
|
||||
for _ in range(3): out = decode(q, kv).numpy()
|
||||
assert out is not None
|
||||
q_ref = q_np[0, :, 0].astype(np.float32)
|
||||
kv_ref = kv_np.astype(np.float32) * scale_np[..., None] if quantized else kv_np.astype(np.float32)
|
||||
k_ref, v_ref = kv_ref[:, 0, :, :valid_kv_len]
|
||||
expected = np.empty((n_heads, 256), dtype=np.float32)
|
||||
for head in range(n_heads):
|
||||
scores = q_ref[head] @ k_ref[head // (n_heads // n_kv_heads)].T / np.sqrt(256)
|
||||
probs = np.exp(scores - scores.max())
|
||||
expected[head] = probs @ v_ref[head // (n_heads // n_kv_heads)] / probs.sum()
|
||||
|
||||
self.assertTrue(np.isfinite(out).all())
|
||||
np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3)
|
||||
|
||||
def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25)
|
||||
|
||||
def test_q8_cache_matches_dequantized_reference(self): self._test_decode(8192, 25, quantized=True)
|
||||
|
||||
def test_q8_cached_prefill_matches_dequantized_reference(self):
|
||||
rng = np.random.default_rng(2)
|
||||
heads, kv_heads, tokens, dim = 16, 2, 32, 256
|
||||
q = rng.standard_normal((1, heads, tokens, dim)).astype(np.float16)
|
||||
kv = rng.standard_normal((2, 1, kv_heads, tokens, dim)).astype(np.float16)
|
||||
scale = np.maximum(np.max(np.abs(kv.astype(np.float32)), axis=-1), 1e-8) / 127
|
||||
packed = np.clip(np.rint(kv.astype(np.float32) / scale[..., None]), -127, 127).astype(np.int8)
|
||||
got = flash_attention_causal_cached(Tensor(q).realize(), Tensor(packed).realize(), tokens, tokens,
|
||||
Tensor(scale.astype(np.float16)).realize()).numpy()
|
||||
dequant = packed.astype(np.float32) * scale.astype(np.float16).astype(np.float32)[..., None]
|
||||
expected = np.empty_like(got)
|
||||
for head in range(heads):
|
||||
scores = q[0, head].astype(np.float32) @ dequant[0, 0, head // (heads // kv_heads)].T / np.sqrt(dim)
|
||||
scores[np.triu_indices(tokens, 1)] = -np.inf
|
||||
probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
|
||||
expected[0, head] = probs @ dequant[1, 0, head // (heads // kv_heads)] / probs.sum(axis=-1, keepdims=True)
|
||||
np.testing.assert_allclose(got, expected, rtol=2e-3, atol=2e-3)
|
||||
|
||||
def test_six_query_heads_per_kv_head(self): self._test_decode(8192, 25, n_heads=12, n_kv_heads=2)
|
||||
|
||||
def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097)
|
||||
|
||||
|
||||
if __name__ == "__main__": unittest.main()
|
||||
|
||||
@@ -1,26 +1,9 @@
|
||||
import unittest, array, time
|
||||
from tinygrad.helpers import mv_address
|
||||
from tinygrad.runtime.support.hcq import MMIOInterface
|
||||
from tinygrad.runtime.support.memory import VirtMapping
|
||||
from tinygrad.runtime.support.system import PCIIfaceBase
|
||||
from tinygrad.runtime.support.usb import USBMMIOInterface
|
||||
from test.mockgpu.usb import MockUSB
|
||||
|
||||
class TestPCIIface(unittest.TestCase):
|
||||
def test_sysmem_mapping_respects_uncached(self):
|
||||
class MM:
|
||||
def alloc_vaddr(self, size, align): return 0x10000
|
||||
def map_range(self, vaddr, size, paddrs, aspace, uncached=False, snooped=False):
|
||||
return VirtMapping(vaddr, size, paddrs, aspace, uncached, snooped)
|
||||
class PCI:
|
||||
def bar_info(self, bar): return 0, 256 << 20
|
||||
def alloc_sysmem(self, size, **kwargs): return memoryview(bytearray(size)), [0x20000]
|
||||
iface = PCIIfaceBase.__new__(PCIIfaceBase)
|
||||
iface.dev, iface.vram_bar, iface.pci_dev = None, 0, PCI()
|
||||
iface.dev_impl = type("DevImpl", (), {"mm": MM()})()
|
||||
for uncached in (False, True):
|
||||
with self.subTest(uncached=uncached): self.assertEqual(iface.alloc(4096, host=True, uncached=uncached).meta.mapping.uncached, uncached)
|
||||
|
||||
class TestHCQIface(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.size = 4 << 10
|
||||
|
||||
@@ -1,105 +0,0 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
|
||||
from tinygrad import Device, Tensor, TinyJit, dtypes
|
||||
from tinygrad.llm.gguf import _GGML_QUANT, ggml_data_to_tensor
|
||||
from tinygrad.llm.kernels import amd as llm_amd
|
||||
from tinygrad.llm.model import Embedding, Linear
|
||||
|
||||
def q8_activation(x:np.ndarray) -> np.ndarray:
|
||||
grouped = x.reshape(*x.shape[:-1], -1, 32)
|
||||
scale = np.maximum(np.max(np.abs(grouped), axis=-1, keepdims=True) / 127, 1e-8)
|
||||
return (np.clip(np.rint(grouped / scale), -127, 127) * scale).reshape(x.shape)
|
||||
|
||||
|
||||
def random_packed(rng:np.random.Generator, ggml_type:int, elements:int) -> np.ndarray:
|
||||
block_size, type_size = _GGML_QUANT[ggml_type]
|
||||
blocks = rng.integers(0, 256, size=(elements // block_size, type_size), dtype=np.uint8)
|
||||
scales = rng.uniform(0.001, 0.02, size=len(blocks)).astype(np.float16).view(np.uint8).reshape(-1, 2)
|
||||
blocks[:, :2] = scales
|
||||
if ggml_type in (12, 13): blocks[:, 2:4] = scales
|
||||
if ggml_type == 14: blocks[:, -2:] = scales
|
||||
return blocks.flatten()
|
||||
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires DEV=AMD")
|
||||
class TestLLMQuantAMD(unittest.TestCase):
|
||||
@staticmethod
|
||||
def assert_q8_equal(result:tuple[Tensor, Tensor, Tensor], expected:np.ndarray):
|
||||
grouped = expected.reshape(expected.shape[0], -1, 32)
|
||||
scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
|
||||
quant = np.clip(np.rint(grouped / scale[..., None]), -127, 127).astype(np.int8)
|
||||
np.testing.assert_equal(result[0].numpy().view(np.int8).reshape(grouped.shape), quant)
|
||||
np.testing.assert_allclose(result[1].numpy(), scale, rtol=1e-6, atol=1e-8)
|
||||
np.testing.assert_equal(result[2].numpy(), quant.astype(np.int32).sum(-1))
|
||||
|
||||
def test_gated_delta_prefill_matches_sequential_reference(self):
|
||||
rng = np.random.default_rng(36)
|
||||
batch, heads, tokens, dim = 1, 2, 5, 128
|
||||
q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
|
||||
q = q / np.linalg.norm(q, axis=-1, keepdims=True) / np.float32(np.sqrt(dim))
|
||||
k = k / np.linalg.norm(k, axis=-1, keepdims=True)
|
||||
beta, alpha = rng.random((batch, heads, tokens), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads, tokens)).astype(np.float32)
|
||||
state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
|
||||
expected_core, expected_state = np.empty_like(q), state.astype(np.float32)
|
||||
for token in range(tokens):
|
||||
state_k = np.einsum("bhij,bhj->bhi", expected_state, k[:, :, token])
|
||||
state_q = np.einsum("bhij,bhj->bhi", expected_state, q[:, :, token])
|
||||
delta = (v[:, :, token] - state_k * alpha[:, :, token, None]) * beta[:, :, token, None]
|
||||
expected_core[:, :, token] = state_q * alpha[:, :, token, None] + delta * np.sum(k[:, :, token] * q[:, :, token], axis=-1)[..., None]
|
||||
expected_state = expected_state * alpha[:, :, token, None, None] + delta[..., None] * k[:, :, token, None, :]
|
||||
core, next_state = llm_amd.gated_delta_prefill(
|
||||
*(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
|
||||
np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
|
||||
np.testing.assert_allclose(next_state.numpy(), expected_state.astype(np.float16), rtol=2e-4, atol=1e-3)
|
||||
|
||||
def test_q8_quantize_matches_reference(self):
|
||||
rng, tokens, in_features = np.random.default_rng(35), 3, 256
|
||||
x = rng.standard_normal((tokens, in_features), dtype=np.float32)
|
||||
grouped = x.reshape(tokens, -1, 32)
|
||||
expected_scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
|
||||
expected_quant = np.clip(np.rint(grouped / expected_scale[..., None]), -127, 127).astype(np.int8)
|
||||
quant, scale, group_sum = llm_amd.q8_quantize_sum(Tensor(x, device="AMD"), tokens, in_features)
|
||||
np.testing.assert_equal(quant.numpy().view(np.int8).reshape(grouped.shape), expected_quant)
|
||||
np.testing.assert_allclose(scale.numpy(), expected_scale, rtol=1e-7, atol=0)
|
||||
np.testing.assert_equal(group_sum.numpy(), expected_quant.astype(np.int32).sum(-1))
|
||||
|
||||
def test_q4_embedding_matches_reference(self):
|
||||
rng, vocab_size, embed_size = np.random.default_rng(34), 16, 256
|
||||
raw = random_packed(rng, 12, vocab_size * embed_size)
|
||||
expected = ggml_data_to_tensor(Tensor(raw), vocab_size * embed_size, 12).reshape(vocab_size, embed_size).half()
|
||||
storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
|
||||
embedding = Embedding(vocab_size, embed_size)
|
||||
embedding.set_quantized(storage[68:], 12)
|
||||
idx = np.array([[7, 1, 15], [0, 4, 7]], dtype=np.int32)
|
||||
np.testing.assert_equal(embedding(Tensor(idx, device="AMD")).numpy(), expected.numpy()[idx])
|
||||
|
||||
def test_iq4_lut_is_ready_for_jit_capture(self):
|
||||
rng, in_features, out_features = np.random.default_rng(33), 256, 16
|
||||
raw = random_packed(rng, 23, out_features * in_features)
|
||||
weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 23).numpy().reshape(out_features, in_features)
|
||||
llm_amd.iq4_half_lut.cache_clear()
|
||||
layer = Linear(in_features, out_features, bias=False)
|
||||
layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 23)
|
||||
@TinyJit
|
||||
def run(x:Tensor): return layer(x).realize()
|
||||
x = rng.standard_normal((16, in_features), dtype=np.float32)
|
||||
expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T
|
||||
for _ in range(2): np.testing.assert_allclose(run(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
|
||||
|
||||
def test_packed_linear_offset_matches_reference(self):
|
||||
rng = np.random.default_rng(32)
|
||||
for ggml_type,in_features in ((8, 256), (12, 256), (13, 256), (14, 256), (23, 256)):
|
||||
for tokens in ((1, 16, 32, 64, 128) if ggml_type == 23 else (1, 16, 128) if ggml_type in (12, 13) else
|
||||
(1, 16) if ggml_type == 14 else (1,)):
|
||||
raw, out_features = random_packed(rng, ggml_type, 64 * in_features), 64
|
||||
weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
|
||||
storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
|
||||
layer = Linear(in_features, out_features, bias=False)
|
||||
layer.set_quantized(storage[68:], ggml_type)
|
||||
x = rng.standard_normal((tokens, in_features), dtype=np.float32)
|
||||
expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T \
|
||||
if ggml_type in (12, 13, 23) and tokens > 1 else q8_activation(x) @ weight.T
|
||||
np.testing.assert_allclose(layer(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
|
||||
|
||||
|
||||
@@ -478,8 +478,15 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
|
||||
return prg
|
||||
|
||||
to_program_cache: dict[tuple, UOp] = {}
|
||||
def to_program(ast:UOp, renderer:Renderer) -> UOp:
|
||||
def program_cache_key(ast:UOp, renderer:Renderer) -> tuple:
|
||||
config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT)
|
||||
key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
|
||||
return (ast.key, type(renderer), renderer.target, *[x.value for x in config])
|
||||
|
||||
def parallel_to_program(args:tuple[UOp, Renderer, tuple]) -> tuple[tuple, UOp]:
|
||||
ast, renderer, key = args
|
||||
return key, do_to_program(ast, renderer)
|
||||
|
||||
def to_program(ast:UOp, renderer:Renderer) -> UOp:
|
||||
key = program_cache_key(ast, renderer)
|
||||
if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
|
||||
return prg
|
||||
|
||||
@@ -1,13 +1,15 @@
|
||||
from __future__ import annotations
|
||||
from typing import cast, Iterator, Any, Sequence
|
||||
import time, random, itertools, math, contextlib, weakref, array
|
||||
import time, random, itertools, math, contextlib, weakref, array, multiprocessing
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
from dataclasses import dataclass, replace, field
|
||||
from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
|
||||
from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
|
||||
from tinygrad.helpers import PARALLEL_COMPILE, NUM_CPU_THREADS
|
||||
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
|
||||
from tinygrad.device import Device, Buffer, MultiBuffer
|
||||
from tinygrad.renderer import Estimates
|
||||
from tinygrad.codegen import to_program
|
||||
from tinygrad.codegen import to_program, to_program_cache, program_cache_key, parallel_to_program
|
||||
from tinygrad.codegen.opt.postrange import bufs_from_ast
|
||||
|
||||
# **************** Helpers ****************
|
||||
@@ -268,6 +270,17 @@ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_li
|
||||
def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp:
|
||||
if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
|
||||
if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
|
||||
if jit and PARALLEL_COMPILE.value:
|
||||
pending:dict[tuple, tuple[UOp, Any, tuple]] = {}
|
||||
for call in linear.toposort():
|
||||
if call.op is not Ops.CALL or call.src[0].op not in (Ops.SINK, Ops.PROGRAM): continue
|
||||
renderer = Device[call.device if isinstance(call.device, str) else call.device[0]].renderer
|
||||
key = program_cache_key(call.src[0], renderer)
|
||||
if key not in to_program_cache: pending.setdefault(key, (call.src[0], renderer, key))
|
||||
if len(pending) >= 16:
|
||||
workers = min(PARALLEL_COMPILE.value, NUM_CPU_THREADS.value, len(pending))
|
||||
with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context("spawn")) as pool:
|
||||
for key,program in pool.map(parallel_to_program, pending.values()): to_program_cache[key] = program
|
||||
linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
|
||||
if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, jit=jit)
|
||||
return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
|
||||
|
||||
@@ -232,6 +232,7 @@ class _DEV(ContextVar):
|
||||
DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
|
||||
IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
|
||||
JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
|
||||
PARALLEL_COMPILE = ContextVar("PARALLEL_COMPILE", 0)
|
||||
CHUNK_SIZE = 2**20 # TinyFS content-addressed store: blob chunk + hash-tree node granularity
|
||||
WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
|
||||
TRAINING = ContextVar("TRAINING", 0)
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Custom kernels used by tinygrad.llm."""
|
||||
+21
-118
@@ -523,112 +523,26 @@ def q4_embedding(layer:Embedding, idx:Tensor) -> Tensor:
|
||||
kernel = _q4_embedding_kernel(params[0], params[1], params[2], params[3], layer.embed_size).call(*srcs)
|
||||
return Tensor(srcs[0].after(kernel))
|
||||
|
||||
@functools.cache
|
||||
def _q8_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, raw_offset:int|UOp=0) -> UOp:
|
||||
if isinstance(raw_offset, UOp): raw_offset = raw_offset.cast(dtypes.uint64)
|
||||
token_tile = 8 if out.shape[0] % 8 == 0 else 1
|
||||
wave_count = 4 if token_tile == 1 else 1
|
||||
token_block, output_block = UOp.range(out.shape[0] // token_tile, 0), UOp.range(out_features // wave_count, 1)
|
||||
wave = UOp.range(wave_count, 3, axis_type=AxisType.LOCAL) if wave_count > 1 else UOp.const(0, dtypes.weakint)
|
||||
output = output_block * wave_count + wave
|
||||
tokens = tuple(token_block * token_tile + i for i in range(token_tile))
|
||||
group_count, lane_count = in_features // 32, min(32, in_features // 32)
|
||||
lane = UOp.range(lane_count, 2, axis_type=AxisType.LOCAL)
|
||||
|
||||
def group_dot(group:UOp) -> list[UOp]:
|
||||
block = output * group_count + group
|
||||
base, odd = raw_offset + block * 8 + block // 2, (block & 1).ne(0)
|
||||
dots = [UOp.const(0, dtypes.int32)] * token_tile
|
||||
for word_idx in range(8):
|
||||
# Q8_0 blocks are 34 bytes, so their two-byte scale makes alternate blocks word-aligned. Read aligned u32s
|
||||
# directly; the other blocks need only two adjacent words instead of four individual byte loads.
|
||||
word = odd.where(raw[base + 1 + word_idx], (raw[base + word_idx] >> 16) | (raw[base + 1 + word_idx] << 16))
|
||||
dots = [_amd_dp4a(word, xq[token, group, word_idx], dot) for token,dot in zip(tokens, dots)]
|
||||
dbits = odd.where(raw[base] >> 16, raw[base] & 0xffff).cast(dtypes.uint16)
|
||||
return [dot.float() * xd[token, group] * dbits.bitcast(dtypes.float16).float() for token,dot in zip(tokens, dots)]
|
||||
|
||||
values = [UOp.const(0, dtypes.float32)] * token_tile
|
||||
for offset in range(0, group_count, lane_count):
|
||||
dots = group_dot((lane + offset).valid(lane + offset < group_count))
|
||||
values = [value + dot for value,dot in zip(values, dots)]
|
||||
totals = [_amd_wave_sum(value, lane, lane_count, wave if wave_count > 1 else None) for value in values]
|
||||
stores = [out[token.valid(lane.eq(0)), output].store(total.cast(out.dtype)) for token,total in zip(tokens, totals)]
|
||||
return UOp.group(*stores).end(token_block, output_block, lane, wave).sink(
|
||||
arg=KernelInfo(name="linear_q8", opts_to_apply=()))
|
||||
|
||||
@functools.cache
|
||||
def _q8_linear_wmma_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, raw_offset:UOp) -> UOp:
|
||||
raw_offset = raw_offset.cast(dtypes.uint64)
|
||||
def load_word(byte_offset:UOp) -> UOp:
|
||||
word_index, half_aligned = byte_offset // 4, (byte_offset & 2).ne(0)
|
||||
word = raw[word_index]
|
||||
return half_aligned.where((word >> 16) | (raw[word_index + 1] << 16), word)
|
||||
token_tile = 32 if out.shape[0] % 32 == 0 else 16
|
||||
token_block, output_block = UOp.range(out.shape[0] // token_tile, 0), UOp.range(out_features // 16, 1)
|
||||
lane = UOp.range(32, 2, axis_type=AxisType.LOCAL)
|
||||
# The codegen may factor this range into several local dimensions. Read the physical wave lane directly.
|
||||
hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
|
||||
physical_col, physical_half = hw_lane % 16, hw_lane // 16
|
||||
output = output_block * 16 + physical_col
|
||||
input_tokens = tuple(token_block * token_tile + tile * 16 + physical_col for tile in range(token_tile // 16))
|
||||
tokens = tuple(tuple(token_block * token_tile + tile * 16 + physical_half * 8 + i for i in range(8))
|
||||
for tile in range(token_tile // 16))
|
||||
group_count = in_features // 32
|
||||
accs = tuple(UOp.placeholder((8,), dtypes.float32, slot=tile, addrspace=AddrSpace.REG) for tile in range(token_tile // 16))
|
||||
accs = tuple(acc.after(acc.store(acc.const_like(0))) for acc in accs)
|
||||
group = UOp.range(group_count, 3, AxisType.REDUCE)
|
||||
raw_accs = [UOp.const(0, dtypes.int32).broadcast(8) for _ in accs]
|
||||
for half in range(2):
|
||||
# rocWMMA's gfx11 loader gives each lane eight values, then appends the eight from lane^16.
|
||||
kbase = half * 16 + physical_half * 8
|
||||
def fragment(words:tuple[UOp, ...]|list[UOp]) -> UOp:
|
||||
swapped_words = tuple(UOp(Ops.CUSTOM, dtypes.uint32, (word,), arg="__builtin_amdgcn_ds_swizzle({0}, 16415)") for word in words)
|
||||
return UOp.stack(*(((word >> (byte * 8)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8)
|
||||
for word in (*words, *swapped_words) for byte in range(4)))
|
||||
block = output * group_count + group
|
||||
bwords = [load_word(raw_offset + block * 34 + 2 + kbase + word * 4) for word in range(2)]
|
||||
bfrag = fragment(bwords)
|
||||
for tile,input_token in enumerate(input_tokens):
|
||||
awords = tuple(xq[input_token, group, kbase // 4 + i].load() for i in range(2))
|
||||
raw_accs[tile] = UOp.wmma(fragment(awords), bfrag, raw_accs[tile], (16, 16, 16), 'AMD', 32)
|
||||
logical_values = []
|
||||
for raw_acc in raw_accs:
|
||||
vals = tuple(raw_acc[i] for i in range(8))
|
||||
swapped = tuple(UOp(Ops.CUSTOM, dtypes.int32, (value,), arg="__builtin_amdgcn_ds_swizzle({0}, 50688)") for value in vals)
|
||||
low = physical_half.eq(0)
|
||||
logical_values.append((low.where(vals[0], swapped[4]), low.where(swapped[0], vals[4]),
|
||||
low.where(vals[1], swapped[5]), low.where(swapped[1], vals[5]),
|
||||
low.where(vals[2], swapped[6]), low.where(swapped[2], vals[6]),
|
||||
low.where(vals[3], swapped[7]), low.where(swapped[3], vals[7])))
|
||||
block = output * group_count + group
|
||||
scale = (load_word(raw_offset + block * 34) & 0xffff).cast(dtypes.uint16).bitcast(dtypes.float16).float()
|
||||
update = UOp.group(*(acc.after(group)[i].store(acc.after(group)[i] + value.float() * scale * xd[token, group])
|
||||
for acc,tile_tokens,logical in zip(accs, tokens, logical_values)
|
||||
for i,(token,value) in enumerate(zip(tile_tokens, logical)))).end(group)
|
||||
stores = [out[token, output].store(acc.after(update)[i]) for acc,tile_tokens in zip(accs, tokens) for i,token in enumerate(tile_tokens)]
|
||||
return UOp.group(*stores).end(token_block, output_block, lane).sink(arg=KernelInfo(name="linear_q8_wmma", opts_to_apply=()))
|
||||
|
||||
@functools.cache
|
||||
def _qk_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xsum:UOp, out_features:int, in_features:int,
|
||||
ggml_type:int, raw_offset:int|UOp=0) -> UOp:
|
||||
raw_offset:int|UOp=0) -> UOp:
|
||||
if isinstance(raw_offset, UOp): raw_offset = raw_offset.cast(dtypes.uint64)
|
||||
def load_byte(base:UOp, byte_offset:UOp) -> UOp:
|
||||
return (raw[base + byte_offset // 4] >> ((byte_offset & 3) * 8).cast(dtypes.uint32)) & 255
|
||||
output_tile = 1
|
||||
output_block, lane = UOp.range(out_features // output_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
|
||||
outputs, group_count = tuple(output_block * output_tile + i for i in range(output_tile)), in_features // 32
|
||||
type_words, output_words = _GGML_QUANT[ggml_type][1] // 4, in_features // 256 * _GGML_QUANT[ggml_type][1] // 4
|
||||
type_words, output_words = _GGML_QUANT[13][1] // 4, in_features // 256 * _GGML_QUANT[13][1] // 4
|
||||
|
||||
def group_dot(group:UOp, output:UOp) -> UOp:
|
||||
block, subgroup = group // 8, group % 8
|
||||
base = raw_offset + output * output_words + block * type_words
|
||||
qs_base = base + (12 if ggml_type == 13 else 4) + (subgroup // 2) * 8
|
||||
qs_base = base + 12 + (subgroup // 2) * 8
|
||||
xwords = _amd_vector_load(xq[0, group, 0], 8)
|
||||
dot = UOp.const(0, dtypes.int32)
|
||||
for word_idx in range(8):
|
||||
word = (raw[qs_base + word_idx] >> ((subgroup & 1) * 4).cast(dtypes.uint32)) & 0x0f0f0f0f
|
||||
if ggml_type == 13:
|
||||
word = word | (((raw[base + 4 + word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4)
|
||||
word = word | (((raw[base + 4 + word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4)
|
||||
dot = _amd_dp4a(word, xwords[word_idx], dot)
|
||||
scale = (subgroup < 4).where(load_byte(base, 4 + subgroup) & 63,
|
||||
(load_byte(base, 8 + subgroup) & 15) | ((load_byte(base, subgroup) >> 6) << 4))
|
||||
@@ -646,11 +560,11 @@ def _qk_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xsum:UOp, out_features:i
|
||||
for acc,output in zip(accs, outputs))).end(chunk)
|
||||
totals = [_amd_wave_sum(acc.after(update)[0].load(), lane, 32) for acc in accs]
|
||||
stores = [out[0, output.valid(lane.eq(0))].store(total.cast(out.dtype)) for output,total in zip(outputs, totals)]
|
||||
return UOp.group(*stores).end(output_block, lane).sink(arg=KernelInfo(name=f"linear_q{ggml_type}", opts_to_apply=()))
|
||||
return UOp.group(*stores).end(output_block, lane).sink(arg=KernelInfo(name="linear_q5_k", opts_to_apply=()))
|
||||
|
||||
@functools.cache
|
||||
def _qk_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int,
|
||||
ggml_type:int, raw_offset:UOp) -> UOp:
|
||||
raw_offset:UOp) -> UOp:
|
||||
x = x.reshape(out.shape[0], in_features)
|
||||
raw_offset = raw_offset.cast(dtypes.uint64)
|
||||
def load_byte(base:UOp, byte_offset:UOp) -> UOp:
|
||||
@@ -667,7 +581,7 @@ def _qk_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_fea
|
||||
input_tokens = tuple(token_block*token_tile + tile*16 + physical_col for tile in range(token_tile // 16))
|
||||
tokens = tuple(tuple(token_block*token_tile + tile*16 + physical_half*8 + i for i in range(8))
|
||||
for tile in range(token_tile // 16))
|
||||
group_count, type_words = in_features // 32, _GGML_QUANT[ggml_type][1] // 4
|
||||
group_count, type_words = in_features // 32, _GGML_QUANT[13][1] // 4
|
||||
output_words = in_features // 256 * type_words
|
||||
accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=output_tile*(token_tile//16)+tile, addrspace=AddrSpace.REG)
|
||||
for tile in range(token_tile // 16)) for output_tile in range(output_tiles))
|
||||
@@ -688,11 +602,10 @@ def _qk_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_fea
|
||||
d = (scales & 0xffff).cast(dtypes.uint16).bitcast(dtypes.float16).float()
|
||||
dmin = (scales >> 16).cast(dtypes.uint16).bitcast(dtypes.float16).float()
|
||||
weight_scale, weight_min = d*scale, dmin*minimum
|
||||
qs_base = base + (12 if ggml_type == 13 else 4) + (subgroup // 2)*8 + half*4
|
||||
qs_base = base + 12 + (subgroup // 2)*8 + half*4
|
||||
qwords = [((raw[qs_base + i] >> ((subgroup & 1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f) for i in range(4)]
|
||||
if ggml_type == 13:
|
||||
qwords = [word | (((raw[base + 4 + half*4 + i] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4)
|
||||
for i,word in enumerate(qwords)]
|
||||
qwords = [word | (((raw[base + 4 + half*4 + i] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4)
|
||||
for i,word in enumerate(qwords)]
|
||||
bfrag = UOp.stack(*(((word >> (byte*8) & 255).float()*weight_scale-weight_min).cast(dtypes.float16)
|
||||
for word in qwords for byte in range(4)))
|
||||
for tile,afrag in enumerate(afrags):
|
||||
@@ -716,7 +629,7 @@ def _qk_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_fea
|
||||
stores = [out[token, output].store(value) for output,output_values in zip(outputs, logical_values)
|
||||
for tile_tokens,logical in zip(tokens, output_values) for token,value in zip(tile_tokens, logical)]
|
||||
return UOp.group(*stores).end(token_block, output_block, lane, wave).sink(
|
||||
arg=KernelInfo(name=f"linear_q{ggml_type}_f16_wmma", opts_to_apply=()))
|
||||
arg=KernelInfo(name="linear_q5_k_f16_wmma", opts_to_apply=()))
|
||||
|
||||
@functools.cache
|
||||
def _iq4_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int,
|
||||
@@ -942,26 +855,26 @@ def _q6_linear_wmma_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, i
|
||||
|
||||
def q8_linear(layer:Linear, x:Tensor, prepared:tuple[Tensor, ...]|None=None) -> Tensor:
|
||||
tokens = int(x.numel()) // layer.in_features
|
||||
if layer.ggml_type in (12, 13):
|
||||
if layer.ggml_type == 13:
|
||||
xq, xd, xsum = prepared if prepared is not None and len(prepared) == 3 else q8_quantize_sum(x, tokens, layer.in_features)
|
||||
else:
|
||||
xq, xd = prepared[:2] if prepared is not None else q8_quantize(x, tokens, layer.in_features)
|
||||
out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
|
||||
if layer._raw_uop is None: layer._prepare_packed()
|
||||
assert layer._raw_uop is not None and layer._raw_offset_uop is not None
|
||||
if layer.ggml_type in (12, 13):
|
||||
if layer.ggml_type == 13:
|
||||
raw_words = layer._raw_uop.bitcast(dtypes.uint32)
|
||||
if tokens % 16 == 0 and layer.out_features % 16 == 0:
|
||||
qk_srcs = (out.uop, raw_words, x.cast(dtypes.float16).contiguous().uop, layer._raw_offset_uop)
|
||||
params = [UOp.placeholder_like(src, slot=i) for i,src in enumerate(qk_srcs)]
|
||||
kernel = _qk_linear_f16_wmma_kernel(params[0], params[1], params[2], layer.out_features,
|
||||
layer.in_features, layer.ggml_type, params[3][0]).call(*qk_srcs)
|
||||
layer.in_features, params[3][0]).call(*qk_srcs)
|
||||
out = Tensor(qk_srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
|
||||
return out if layer.bias is None else out + layer.bias
|
||||
qk_decode_srcs = (out.uop, raw_words, xq.uop, xd.uop, xsum.uop, layer._raw_offset_uop)
|
||||
params = [UOp.placeholder_like(src, slot=i) for i,src in enumerate(qk_decode_srcs)]
|
||||
kernel = _qk_linear_kernel(params[0], params[1], params[2], params[3], params[4], layer.out_features,
|
||||
layer.in_features, layer.ggml_type, params[5][0]).call(*qk_decode_srcs)
|
||||
layer.in_features, params[5][0]).call(*qk_decode_srcs)
|
||||
out = Tensor(qk_decode_srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
|
||||
return out if layer.bias is None else out + layer.bias
|
||||
if layer.ggml_type == 23:
|
||||
@@ -978,22 +891,13 @@ def q8_linear(layer:Linear, x:Tensor, prepared:tuple[Tensor, ...]|None=None) ->
|
||||
layer.in_features, params[4][0])).call(*iq4_srcs)
|
||||
out = Tensor(iq4_srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
|
||||
return out if layer.bias is None else out + layer.bias
|
||||
raw = layer._raw_uop.bitcast(dtypes.uint32) if layer.ggml_type == 8 else layer._raw_uop
|
||||
srcs = (out.uop, raw, xq.uop, xd.uop, layer._raw_offset_uop)
|
||||
assert layer.ggml_type == 14
|
||||
srcs = (out.uop, layer._raw_uop, xq.uop, xd.uop, layer._raw_offset_uop)
|
||||
params = [UOp.placeholder_like(src, slot=i) for i,src in enumerate(srcs)]
|
||||
if layer.ggml_type == 8:
|
||||
if tokens % 16 == 0 and layer.out_features % 16 == 0:
|
||||
kernel = _q8_linear_wmma_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0] * 4).call(*srcs)
|
||||
else:
|
||||
kernel = _q8_linear_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0]).call(*srcs)
|
||||
else:
|
||||
assert layer.ggml_type == 14
|
||||
kernel = (_q6_linear_wmma_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0] * 4) if tokens % 16 == 0 and layer.out_features % 16 == 0 else
|
||||
_q6_linear_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0] * 4)).call(*srcs)
|
||||
kernel = (_q6_linear_wmma_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0] * 4) if tokens % 16 == 0 and layer.out_features % 16 == 0 else
|
||||
_q6_linear_kernel(params[0], params[1], params[2], params[3], layer.out_features,
|
||||
layer.in_features, params[4][0] * 4)).call(*srcs)
|
||||
out = Tensor(srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
|
||||
return out if layer.bias is None else out + layer.bias
|
||||
|
||||
@@ -1002,4 +906,3 @@ def iq4_half_lut(device:str) -> Tensor:
|
||||
from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
|
||||
values = [x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])]
|
||||
return Tensor(values, dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous().realize()
|
||||
|
||||
|
||||
@@ -42,13 +42,13 @@ class Linear(nn.Linear, PackedWeight):
|
||||
self.in_features, self.out_features = in_features, out_features
|
||||
self._init_packed()
|
||||
def prepare(self, x:Tensor, with_sum:bool=False) -> tuple[Tensor, ...]|None:
|
||||
if (with_sum or self.ggml_type in (12, 13)) and self.ggml_type in (8, 12, 13, 14, 23) and \
|
||||
if (with_sum or self.ggml_type == 13) and self.ggml_type in (13, 14, 23) and \
|
||||
str(self.weight.device).startswith("AMD"):
|
||||
return llm_amd.q8_quantize_sum(x, int(x.numel()) // self.in_features, self.in_features)
|
||||
return llm_amd.q8_quantize(x, int(x.numel()) // self.in_features, self.in_features) \
|
||||
if self.ggml_type in (8, 14, 23) and str(self.weight.device).startswith("AMD") else None
|
||||
if self.ggml_type in (14, 23) and str(self.weight.device).startswith("AMD") else None
|
||||
def __call__(self, x:Tensor, prepared:tuple[Tensor, ...]|None=None) -> Tensor:
|
||||
if self.ggml_type in (8, 12, 13, 14, 23) and str(self.weight.device).startswith("AMD"):
|
||||
if self.ggml_type in (13, 14, 23) and str(self.weight.device).startswith("AMD"):
|
||||
return llm_amd.q8_linear(self, x, prepared)
|
||||
if self.ggml_type is not None:
|
||||
weight = ggml_data_to_tensor(self.weight, self.out_features * self.in_features, self.ggml_type,
|
||||
@@ -653,7 +653,7 @@ class Transformer:
|
||||
quantization = get_ggml_quantization(weight)
|
||||
owner = resolve_owner(parts[:-1]) if parts[-1] == "weight" else None
|
||||
packed = quantization is not None and str(load_device).startswith("AMD") and \
|
||||
(isinstance(owner, Linear) and quantization[1] in (8, 12, 13, 14, 23) or
|
||||
(isinstance(owner, Linear) and quantization[1] in (13, 14, 23) or
|
||||
isinstance(owner, Embedding) and quantization[1] == 12 and str(load_device).startswith("AMD"))
|
||||
if packed:
|
||||
assert quantization is not None and isinstance(owner, PackedWeight)
|
||||
@@ -718,6 +718,7 @@ class Transformer:
|
||||
assert self._restore_state_jit is not None
|
||||
self._restore_state_jit()
|
||||
|
||||
@Context(PARALLEL_COMPILE=getenv("PARALLEL_COMPILE", 12))
|
||||
def warmup(self, chunk_size:int=256):
|
||||
device = self.token_embd.weight.device
|
||||
direct_capture = not self.has_recurrent_block and all(isinstance(block, TransformerBlock) for block in self.blk)
|
||||
|
||||
Reference in New Issue
Block a user