work

test/unit/test_function.py

@@ -557,6 +557,15 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
+  def test_realize_inside_function(self):
+    constants = []
+    @function(precompile=True, allow_implicit=True)
+    def f(a:Tensor):
+      constants.append(Tensor([2], device=a.device).realize())
+      return a + constants[-1]
+    np.testing.assert_equal(f(Tensor([1], device="CPU").realize()).numpy(), [3])
+    self.assertTrue(constants[0].uop.is_realized)
+
   def test_custom_kernel_write_only_persistent_output_is_implicit(self):
     # a write-only custom_kernel output that is a realized buffer must be captured
     def write(C:UOp, A:UOp) -> UOp:

test/unit/test_llm_amd.py

@@ -1,10 +1,9 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes
-from tinygrad.llm.kernels import cached_attention
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.llm.kernels import Linear, cached_attention
 from tinygrad.llm.kernels.amd import q8_quantize
 from tinygrad.llm.gguf import ggml_data_to_tensor
-from tinygrad.llm.model import Linear
 
 class TestQ8Quantize(unittest.TestCase):
   def test_values_and_scales(self):
@@ -23,10 +22,9 @@ class TestQ8Quantize(unittest.TestCase):
     raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
     decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
     linear = Linear(256, 1, bias=False)
-    offset = linear.set_quantized(decoded)
-    assert offset is not None
-    linear._raw_offset_uop = offset.realize().uop
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
     self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
+    self.assertIsNotNone(linear._raw_offset)
 
   def test_attention_uses_physical_cache_length(self):
     if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required")

tinygrad/llm/kernels/__init__.py

@@ -10,7 +10,7 @@ class Linear(nn.Linear):
   def __init__(self, in_features:int, out_features:int, bias=True):
     super().__init__(in_features, out_features, bias)
     self.in_features, self.out_features = in_features, out_features
-    self._raw_offset_uop:UOp|None = None
+    self._raw_offset:Tensor|None = None
   def set_quantized(self, decoded:Tensor) -> Tensor|None:
     packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((13, 176), (14, 210), (23, 136))}
     raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
@@ -21,8 +21,10 @@ class Linear(nn.Linear):
     if self.ggml_type == 23 and str(self.weight.device).startswith("AMD"):
       from tinygrad.llm.kernels.amd import iq4_half_lut
       iq4_half_lut(str(self.weight.device))
-    return Tensor([raw_offset // 4], dtype=dtypes.uint64, device=self.weight.device)
+    self._raw_offset = Tensor([raw_offset//4], dtype=dtypes.uint64, device=self.weight.device).realize()
+    return self._raw_offset
   def __call__(self, x:Tensor) -> Tensor:
+    if self.ggml_type is None and str(self.weight.device).startswith("AMD"): self.set_quantized(self.weight)
     if self.ggml_type in (13, 14, 23) and str(self.weight.device).startswith("AMD"):
       from tinygrad.llm.kernels.amd import q8_linear
       return q8_linear(self, x)

tinygrad/llm/kernels/amd.py

@@ -395,10 +395,10 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp
   return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
 
 def q8_linear(layer:Linear, x:Tensor) -> Tensor:
-  assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_offset_uop is not None
+  assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_offset is not None
   tokens = int(x.numel()) // layer.in_features
   out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop
-  raw, offset = layer.weight.uop.buf_uop, layer._raw_offset_uop
+  raw, offset = layer.weight.uop.buf_uop, layer._raw_offset.uop
   out_features, in_features = layer.out_features, layer.in_features
   use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0
   def run(fxn:Callable[..., UOp], *srcs:UOp) -> Tensor:

tinygrad/llm/model.py

@@ -1,21 +1,11 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from typing import cast
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes
 from tinygrad.llm.kernels import Linear, cached_attention, gated_delta_prefill
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
-def load_state_dict(model:Transformer, state_dict:dict[str, Tensor]):
-  nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)
-  layers = cast(dict[str, Linear], nn.state.get_state_dict(model, tensor_type=Linear)).values()
-  packed:list[tuple[Linear, Tensor]] = []
-  for layer in layers:
-    if str(layer.weight.device).startswith("AMD") and (offset:=layer.set_quantized(layer.weight)) is not None: packed.append((layer, offset))
-  if packed: Tensor.realize(*(offset for _,offset in packed))
-  for layer,offset in packed: layer._raw_offset_uop = offset.uop
-
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
   freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
@@ -408,7 +398,7 @@ class Transformer:
       qkv_bias='blk.0.attn_q.bias' in state_dict,
       expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict)
     model = Transformer(config)
-    load_state_dict(model, state_dict)  # NOTE: rope_freqs.weight (32,) is unused
+    nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     if realize:
       for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())

tinygrad/tensor.py

@@ -4,7 +4,7 @@ import time, functools, sys, inspect, pathlib, hashlib, weakref
 from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
-from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
+from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey, Context
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike
 from tinygrad.mixin.rand import RandMixin
@@ -16,6 +16,16 @@ from tinygrad.callify import transform_to_call
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
+def _realize_tensors(tensors:tuple[Tensor, ...], do_update_stats=True):
+  for x in tensors:
+    if x.uop.op is Ops.COPY and x.uop.src[0].op is Ops.BUFFER and x.uop.src[0].device == "PYTHON" and \
+      isinstance(x.device, str) and x.nbytes() <= 8:
+      out = UOp.new_buffer(x.device, x.uop.max_numel(), x.dtype)
+      with Context(ALLOW_DEVICE_USAGE=1):
+        cast(Buffer, out.buffer).ensure_allocated().copy_from(cast(Buffer, x.uop.src[0].buffer).ensure_allocated())
+      x.uop = out.reshape(x.shape)
+  to_realize = [x for x in tensors if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
+  if to_realize: run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
 def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
@@ -190,9 +200,7 @@ class Tensor(RandMixin):
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
-    to_realize = [x for x in (self,)+lst if not x.uop.is_virtual and not x.uop.has_buffer_identity()]
-    if len(to_realize):
-      run_linear(*Tensor.linear_with_vars(*to_realize), update_stats=do_update_stats)
+    _realize_tensors((self,)+lst, do_update_stats)
     return self
 
   def replace(self, x:Tensor) -> Tensor: