Apple GPU R1 (start): device-resident bmm (op-to-op residency)

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 |

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.
 

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<MTLBuffer> bufA,
+                             id<MTLBuffer> bufB, id<MTLBuffer> 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<NSNumber *> *aShape = @[ @(batch), @(M), @(K) ];
+    NSArray<NSNumber *> *bShape = @[ @(bBatch), @(K), @(N) ];
+    NSArray<NSNumber *> *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,

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<std::size_t>(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<const float*>(A->data),
+                            static_cast<const float*>(B->data),
+                            static_cast<float*>(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,

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