diff --git a/docs/apple_gpu_resident_activations_plan.md b/docs/apple_gpu_resident_activations_plan.md index da0d0f06b..2ae1bde19 100644 --- a/docs/apple_gpu_resident_activations_plan.md +++ b/docs/apple_gpu_resident_activations_plan.md @@ -102,8 +102,8 @@ removes the per-op syncs. | Phase | Scope | Effort | Risk | |---|---|---|---| -| **R0** | `DeviceTensor` handle + buffer registry (`alloc`/`wrap`/`upload`/`download`/`free`); zero-copy `wrap` of numpy via shared buffers; Python `DeviceTensor` + lazy `_as_numpy`. No kernel changes. | Med | Low | -| **R1** | Handle-taking entry points for the hot decode ops (matmul/bmm/absorb_decode/rowops/gumbel). Keep host-ptr variants for compat. Dispatch loop threads handles; materialize lazily. | Med–Large | Med | +| **R0 ✅** | **DONE** — `DeviceTensor` handle + `ts_dev_alloc/contents/nbytes/upload/download/free/is_metal`; shared-buffer storage with zero-copy `.numpy()` view + lazy host materialization; non-Apple host-memory parity. `tests/unit/test_apple_gpu_device_tensor.py` (13). | Med | Low | +| **R1 🟡** | **Started** — first handle-taking entry point: `tessera_apple_gpu_bmm_dev_f32(TsDeviceTensor A,B,O,…)` consumes the inputs' shared buffers in place and writes the MPSGraph result straight into the output buffer (`resultsDictionary`) — **no host upload, no readback** — so a chain of `_apple_gpu_bmm_device` calls keeps intermediates on-GPU. Shares the bmm graph cache; host-ptr path kept. `tests/unit/test_apple_gpu_resident_bmm.py` (6, incl. a 4-deep chain). **Remaining:** extend handle entry points to absorb_decode / rowops / gumbel and thread handles through the metadata dispatch loop with lazy materialization. | Med–Large | Med | | **R2** | **Command-buffer batching** — one `encodeToCommandBuffer:` per op-chain, single commit + wait. The core perf lever. | Large | **High** | | **R3** | Persistent decode-loop state — logits resident → Gumbel sampler consumes a device handle → only the token id reads back. | Med | Med | | **R4** | Device-resident KV cache — `c_kv`/`k_rope` and the paged blocks live in device buffers; gather/append happen on-device. | Large | Med–High | diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index 73f9364d0..bda7fc2fc 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -4971,6 +4971,55 @@ def __repr__(self) -> str: f"{loc}, freed={self._freed})") +_BMM_DEV_CONFIGURED = False + + +def _apple_gpu_bmm_dev_f32() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_bmm_dev_f32", None) + if sym is None: + return None + global _BMM_DEV_CONFIGURED + if not _BMM_DEV_CONFIGURED: + vp, i32 = ctypes.c_void_p, ctypes.c_int32 + sym.argtypes = [vp, vp, vp, i32, i32, i32, i32, i32] + sym.restype = i32 + _BMM_DEV_CONFIGURED = True + return sym + + +def _apple_gpu_bmm_device(A: "DeviceTensor", B: "DeviceTensor", + b_broadcast: bool = False) -> "DeviceTensor | None": + """R1 — device-resident batched matmul. Both inputs are ``DeviceTensor``s + (shapes ``[batch, M, K]`` and ``[batch|1, K, N]``); the result is a new + ``DeviceTensor`` ``[batch, M, N]`` that stays on-device — **no host upload + or readback**, so it can feed the next op directly. f32 only. Returns None + when the device path is unavailable.""" + import numpy as _np + if A.dtype != _np.float32 or B.dtype != _np.float32: + return None + if len(A.shape) != 3 or len(B.shape) != 3: + return None + batch, M, K = A.shape + bBatch, K2, N = B.shape + if K2 != K or (bBatch != batch and bBatch != 1): + return None + bcast = (bBatch == 1 and batch != 1) or bool(b_broadcast) + sym = _apple_gpu_bmm_dev_f32() + if sym is None: + return None + out = DeviceTensor.empty((batch, M, N), _np.float32) + if out is None: + return None + rc = sym(A.handle, B.handle, out.handle, ctypes.c_int32(batch), + ctypes.c_int32(M), ctypes.c_int32(N), ctypes.c_int32(K), + ctypes.c_int32(1 if bcast else 0)) + if rc != 1: + out.free() + return None + return out + + def _load_apple_gpu_runtime() -> ctypes.CDLL: """Phase 8.3: locate or compile the apple_gpu runtime shared library. diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm index 96fa73d4a..15ed2a508 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm @@ -8313,8 +8313,74 @@ static bool mpsg_run_bmm(MetalDeviceContext &ctx, const void *a, const void *b, } } +// R1 — device-resident bmm: inputs + output are existing shared MTLBuffers +// (from DeviceTensor handles), so there is NO host upload and NO readback. The +// MPSGraph result is written straight into the output buffer via +// resultsDictionary, leaving it resident for the next op. Shares the bmm graph +// cache key, so it reuses the same compiled graph as the host-ptr path. +static bool mpsg_run_bmm_dev(MetalDeviceContext &ctx, id bufA, + id bufB, id bufO, + int32_t batch, int32_t M, int32_t N, int32_t K, + bool b_broadcast, MPSDataType ioType) { + if (batch <= 0 || M <= 0 || N <= 0 || K <= 0) return true; + if (!bufA || !bufB || !bufO) return false; + @autoreleasepool { + int32_t bBatch = b_broadcast ? 1 : batch; + NSArray *aShape = @[ @(batch), @(M), @(K) ]; + NSArray *bShape = @[ @(bBatch), @(K), @(N) ]; + NSArray *oShape = @[ @(batch), @(M), @(N) ]; + NSString *key = [NSString stringWithFormat:@"bmm:%d:%d:%d:%d:%d:%d", + (int)ioType, batch, M, N, K, + (int)b_broadcast]; + NSArray *entry = mpsg_cache_get(key); + MPSGraph *g; + MPSGraphTensor *pa, *pb, *y; + if (entry) { + g = entry[0]; + pa = ((NSArray *)entry[1])[0]; + pb = ((NSArray *)entry[1])[1]; + y = entry[2]; + } else { + g = [MPSGraph new]; + pa = [g placeholderWithShape:aShape dataType:ioType name:nil]; + pb = [g placeholderWithShape:bShape dataType:ioType name:nil]; + MPSGraphTensor *yf = + [g matrixMultiplicationWithPrimaryTensor:mpsg_up(g, pa, ioType) + secondaryTensor:mpsg_up(g, pb, ioType) + name:nil]; + y = mpsg_down(g, yf, ioType); + mpsg_cache_put(key, @[ g, @[ pa, pb ], y ]); + } + MPSGraphTensorData *ad = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufA shape:aShape dataType:ioType]; + MPSGraphTensorData *bd = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufB shape:bShape dataType:ioType]; + MPSGraphTensorData *od = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufO shape:oShape dataType:ioType]; + [g runWithMTLCommandQueue:ctx.queue + feeds:@{pa : ad, pb : bd} + targetOperations:nil + resultsDictionary:@{y : od}]; + return true; + } +} + } // namespace +// R1 device-resident bmm entry point. A/B/O are TsDeviceTensor handles whose +// shared buffers are used in place. Returns 1 on a real GPU run, 0 otherwise +// (caller falls back to the host-ptr path). +extern "C" int32_t tessera_apple_gpu_bmm_dev_f32(TsDeviceTensor *A, + TsDeviceTensor *B, + TsDeviceTensor *O, + int32_t batch, int32_t M, + int32_t N, int32_t K, + int32_t b_broadcast) { + MetalDeviceContext &ctx = deviceContext(); + if (!ctx.ok || !A || !B || !O) return 0; + return mpsg_run_bmm_dev(ctx, A->buf, B->buf, O->buf, batch, M, N, K, + b_broadcast != 0, MPSDataTypeFloat32) + ? 1 + : 0; +} + extern "C" void tessera_apple_gpu_bmm_f32(const float *A, const float *B, float *O, int32_t batch, int32_t M, int32_t N, int32_t K, diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp index 16ac3f84b..26309318a 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp @@ -1404,6 +1404,20 @@ extern "C" void tessera_apple_gpu_bmm_f16(const uint16_t*, const uint16_t*, // runtime.py upcasts to f32 on the fallback path. std::memset(O, 0, static_cast(batch) * M * N * 2); } +// R1 device-resident bmm — non-Apple reference. The handles are host-memory +// backed, so this is the same bmm into O->data. +extern "C" int32_t tessera_apple_gpu_bmm_dev_f32(TsDeviceTensor* A, + TsDeviceTensor* B, + TsDeviceTensor* O, int32_t batch, + int32_t M, int32_t N, int32_t K, + int32_t b_broadcast) { + if (!A || !B || !O) return 0; + tessera_apple_gpu_bmm_f32(static_cast(A->data), + static_cast(B->data), + static_cast(O->data), batch, M, N, K, + b_broadcast); + return 1; +} // ---- Tier-3 reduction lane non-Apple reference (2026-05-29) ---------------- extern "C" void tessera_apple_gpu_mpsgraph_reduce_f32(int32_t op, const float* x, diff --git a/tests/unit/test_apple_gpu_resident_bmm.py b/tests/unit/test_apple_gpu_resident_bmm.py new file mode 100644 index 000000000..b81db7e14 --- /dev/null +++ b/tests/unit/test_apple_gpu_resident_bmm.py @@ -0,0 +1,102 @@ +"""Apple GPU R1 — device-resident bmm (op-to-op residency). + +`runtime._apple_gpu_bmm_device` consumes and produces `DeviceTensor` handles: +the inputs' shared buffers are used in place (no host upload) and the MPSGraph +result is written straight into the output buffer (no readback). So a chain of +device-resident ops keeps its intermediates on-GPU — the mechanism that lets the +decode loop stop round-tripping activations. Validated against numpy. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tessera import runtime as R +from tessera.runtime import DeviceTensor + + +def _require(): + if R._apple_gpu_bmm_dev_f32() is None: + pytest.skip("device-resident bmm unavailable") + + +def _ref_bmm(A, B): + if B.shape[0] == 1 and A.shape[0] != 1: + B = np.broadcast_to(B, (A.shape[0], B.shape[1], B.shape[2])) + return np.matmul(A.astype(np.float64), B.astype(np.float64)) + + +def test_device_bmm_matches_numpy(): + _require() + rng = np.random.RandomState(0) + A = rng.randn(3, 4, 5).astype(np.float32) + B = rng.randn(3, 5, 6).astype(np.float32) + da, db = DeviceTensor.from_numpy(A), DeviceTensor.from_numpy(B) + out = R._apple_gpu_bmm_device(da, db) + assert out is not None and out.shape == (3, 4, 6) + np.testing.assert_allclose(out.numpy(), _ref_bmm(A, B), rtol=1e-4, atol=1e-4) + for t in (da, db, out): + t.free() + + +def test_device_bmm_broadcast_B(): + _require() + rng = np.random.RandomState(1) + A = rng.randn(4, 3, 8).astype(np.float32) + B = rng.randn(1, 8, 7).astype(np.float32) # shared across batch + da, db = DeviceTensor.from_numpy(A), DeviceTensor.from_numpy(B) + out = R._apple_gpu_bmm_device(da, db) + assert out is not None and out.shape == (4, 3, 7) + np.testing.assert_allclose(out.numpy(), _ref_bmm(A, B), rtol=1e-4, atol=1e-4) + + +def test_chain_keeps_intermediate_resident(): + """C = bmm(A, B); D = bmm(C, E). The intermediate C is consumed by the + second bmm directly as a DeviceTensor — it is never materialized to host + (we only call .numpy() on the final D).""" + _require() + rng = np.random.RandomState(2) + A = rng.randn(2, 4, 5).astype(np.float32) + B = rng.randn(2, 5, 6).astype(np.float32) + E = rng.randn(2, 6, 3).astype(np.float32) + da, db, de = (DeviceTensor.from_numpy(x) for x in (A, B, E)) + + C = R._apple_gpu_bmm_device(da, db) # resident intermediate + assert C is not None + D = R._apple_gpu_bmm_device(C, de) # consumes C without a readback + assert D is not None and D.shape == (2, 4, 3) + + ref = np.matmul(_ref_bmm(A, B), E.astype(np.float64)) + np.testing.assert_allclose(D.numpy(), ref, rtol=1e-4, atol=1e-4) + for t in (da, db, de, C, D): + t.free() + + +def test_resident_output_feeds_host_only_at_end(): + """A 3-deep chain; only the final output is read back to host.""" + _require() + rng = np.random.RandomState(3) + mats = [rng.randn(1, 8, 8).astype(np.float32) for _ in range(4)] + dts = [DeviceTensor.from_numpy(m) for m in mats] + acc = dts[0] + for nxt in dts[1:]: + acc = R._apple_gpu_bmm_device(acc, nxt) + assert acc is not None + ref = mats[0].astype(np.float64) + for m in mats[1:]: + ref = np.matmul(ref, m.astype(np.float64)) + np.testing.assert_allclose(acc.numpy(), ref, rtol=1e-3, atol=1e-3) + + +def test_device_bmm_rejects_non_f32(): + _require() + a = DeviceTensor.from_numpy(np.zeros((2, 3, 4), np.float16)) + b = DeviceTensor.from_numpy(np.zeros((2, 4, 5), np.float16)) + assert R._apple_gpu_bmm_device(a, b) is None + + +def test_symbol_exported(): + rt = R._load_apple_gpu_runtime() + assert hasattr(rt, "tessera_apple_gpu_bmm_dev_f32") + assert R._apple_gpu_bmm_dev_f32() is not None