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

[gstoner]

First handle-taking kernel entry point in the GPU-resident activation build-out
(docs/apple_gpu_resident_activations_plan.md). Proves the core mechanism:
a producer op's output feeds the consumer with no host round-trip.

What

• tessera_apple_gpu_bmm_dev_f32(TsDeviceTensor A, B, O, …) — the inputs'
  shared MTLBuffers (from DeviceTensor handles, R0) are used in place (no
  newBufferWithBytes upload), and the MPSGraph result is written straight
  into the output buffer via resultsDictionary (no readBytes). It reuses
  the existing bmm graph cache key, so it shares the compiled graph with the
  host-ptr path.
• runtime._apple_gpu_bmm_device(A, B) — both inputs DeviceTensors,
  returns a resident DeviceTensor [batch,M,N]. f32; [batch,M,K] × [batch|1,K,N] (with B-broadcast).
• Non-Apple stub parity (host bmm into O->data).

The proof

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,
never materialized to host; only the final D is read back. A 4-deep chain test
reads back only at the very end. This is the mechanism that lets the decode loop
stop round-tripping activations.

Tests

tests/unit/test_apple_gpu_resident_bmm.py — 6 tests: vs numpy, B-broadcast,
2-bmm resident chain, 4-deep chain (host read only at the end), non-f32
rejection, symbol export.

Scope (this PR vs. the rest of R1)

This lands the keystone (bmm) and proves op-to-op residency end-to-end.
Remaining R1 (follow-up PRs): handle entry points for absorb_decode /
rowops / the Gumbel sampler, and threading handles through the metadata dispatch
loop with lazy materialization.

Verification (local, Apple Silicon)

• resident-bmm + device-tensor suite: 19/19; apple_gpu + mla + buffer-pool sweep: 534 passed
• buffer-pool RAII gate: pass (the dev path uses existing buffers — no raw allocs)
• mypy ratchet: clean (only the pre-existing environmental torch-import error)

CI on this repo is uniformly red on main (Python 3.8–3.11 matrix, missing
optional deps) — same state PRs #17–#33 merged through. The local signal above
is green.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3c0ca14c7c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Require resident-bmm symbol before accepting cached runtime

When TESSERA_APPLE_GPU_RUNTIME_LIB or the CMake build directory points at a library built before this change, _load_apple_gpu_runtime() still accepts it because its symbol gate was not extended for tessera_apple_gpu_bmm_dev_f32; this new lookup then returns None and every resident-bmm call is silently unavailable instead of forcing the rebuild path the loader uses for prior runtime additions. Please add the new symbol to the loader's required-symbol checks so stale cached runtimes do not disable the feature.

Useful? React with 👍 / 👎.