diff --git a/.agents/backend-matrix.md b/.agents/backend-matrix.md index 47432afc0..2b89d5bc7 100644 --- a/.agents/backend-matrix.md +++ b/.agents/backend-matrix.md @@ -246,7 +246,8 @@ this repository. State remains `ACTIVE`; no lifecycle transition is claimed. | `BACKEND-TENSTORRENT` | Tenstorrent Blackhole (Tensix multicore, discrete PCIe, no unified memory) — thin `vt::` adapter over ttnn's existing C++ op library rather than hand-written kernels, mirroring the Metal/MLX decision (E1); vLLM has no Tenstorrent platform anywhere | vllm.cpp extension through upstream seam `platforms/interface.py:134-229` (same pattern as Metal/Vulkan) | **ACTIVE 2026-08-10.** `vt::tenstorrent::Backend` + registrar [tenstorrent_backend.cpp](../src/vt/tenstorrent/tenstorrent_backend.cpp); shared mesh-device lifecycle [tenstorrent_device.cpp](../src/vt/tenstorrent/tenstorrent_device.cpp); 17 registered ops cover OPT-125m and the Qwen3-0.6B forward (`kMatmul`, `kMatmulBT`, `kAdd`, `kRelu`, `kEmbedding`, `kLayerNorm`, `kRmsNorm`, `kSiluAndMul`, bf16/f32 casts, three RoPE forms, `kQkvSplit`, `kReshapeAndCache`, host-oracle `kPagedAttention`, `kGreedyArgmax`) [tenstorrent_ops.cpp](../src/vt/tenstorrent/tenstorrent_ops.cpp); platform allow-list selects OPT and Qwen3 [platforms/tenstorrent.cpp](../src/vllm/platforms/tenstorrent.cpp). `DeviceType::kTENSTORRENT` [device.h](../include/vt/device.h) | [test_tenstorrent_backend.cpp](../tests/vt/test_tenstorrent_backend.cpp) carries real-Blackhole op gates; [test_qwen3_paged_engine.cpp](../tests/parity/test_qwen3_paged_engine.cpp) selects Tenstorrent device-specific anchor and teacher-forced near-tie goldens. OPT-125m STRICT 6/6 passed. Qwen3 short warm smoke ran 4 tokens at about 0.28 tok/s; full 16x16 gate remains pending behind host paged attention | [tenstorrent-backend.md](specs/tenstorrent-backend.md) | `ACTIVE` | `CLAIM-BACKEND-TENSTORRENT-SPIKE` | | `BACKEND-TENSTORRENT-RESIDUAL-GOLDEN` | Child of `BACKEND-TENSTORRENT` — the owed op-level numerics evidence at the residual-RMS device boundary (`kDeviceResidualMinRows == 32`): device path does `ttnn::add`+`ttnn::rms_norm` in bf16; host/CPU path accumulates in f32. Bot-flagged on #289; never measured at the boundary. | vllm.cpp CPU oracle `RmsNormKernel` mirrors vLLM `fused_add_rms_norm` (add in model dtype, variance in f32); `src/vt/cpu/cpu_ops.cpp:371-398` | `src/vt/tenstorrent/tenstorrent_ops.cpp:1067-1117` (host/device split, `kDeviceResidualMinRows=32`) | [test_tenstorrent_backend.cpp](../tests/vt/test_tenstorrent_backend.cpp) `kRmsNorm residual: device vs CPU f32 oracle across the rows=32 boundary`: 22/22 cases on real Blackhole P150. **Measured 2026-08-11:** host path `rows<32` bit-identical to CPU (`max_abs=0`); device bf16 path `rows>=32` diverges by constant **0.0459 abs** (1.9–2.6× rel on near-zero outputs) — bf16 rounding signature, not accumulation. Decision pending the e2e golden tie-break | [tenstorrent-residual-golden.md](specs/tenstorrent-residual-golden.md) | `SPIKE` | `CLAIM-BACKEND-TENSTORRENT-RESIDUAL-GOLDEN` | | `BACKEND-TENSTORRENT-MISTRAL` | Child of `BACKEND-TENSTORRENT` — allowlist `MistralForCausalLM` (Mistral-7B-v0.3: GQA 32/8, head_dim 128, plain rope theta 1e6, untied lm_head, full attention) on the TT platform + device-aware SACRED gate. Mistral reuses the Qwen3-dense forward verbatim (qk-norm skipped); every op already registered. No new kernel. | vLLM `mistral.py::MistralForCausalLM(LlamaForCausalLM)` (already ported to the shared dense machinery); gate pattern mirrored from `test_qwen3_paged_engine.cpp:221-296` | `src/vllm/platforms/tenstorrent.cpp:52-54` (allowlist) + `tests/parity/test_mistral_paged_engine.cpp` (device-aware wiring + Backend Proof) | **Gate PASSED on Blackhole P150 (2026-08-12):** [test_mistral_paged_engine.cpp](../tests/parity/test_mistral_paged_engine.cpp) 16/16 prompts PASS (12/16 strict-exact, 4/16 near-tie, 0 forward-divergent), max gap **0.062 nats**, BACKEND PROOF 0 declines (kMatmul selections=256 = untied lm_head on device, kPagedAttention=8192). Goldens `our_ids_tenstorrent.npy` + `neartie_gap_mnats_tenstorrent.npy` (transformers alternative-oracle; POL-ORACLE deviation recorded, same as Qwen3-0.6B TT precedent). Exit SIGSEGV 139 is the known MeshDevice teardown crash, not a gate failure | [tenstorrent-mistral.md](specs/tenstorrent-mistral.md) | `ACTIVE` | `CLAIM-BACKEND-TENSTORRENT-MISTRAL` | - +| `BACKEND-TENSTORRENT-TRACE-RUNNER` | Child of `BACKEND-TENSTORRENT` — wire the landed graph-capture foundation (#354 / `59568772`) into a capturable forward region. Handoff §8/§9 "not done". **Decision (2026-08-13): NO-GO for pure T=1 decode capture — MEASURED, not assumed.** | CUDA is the ONLY backend with `SupportsGraphCapture()==true` (`cuda_backend.cu:184-240`); Metal/Vulkan `false`. Shared decode-graph framework `Qwen3DenseDecodeGraph` (qwen3.cpp:489, used by Qwen3/Mistral/Llama/InternLM2) gated on `support_static_graph_mode()` which TT does NOT override (base `false`) | `src/vt/tenstorrent/tenstorrent_backend.cpp:70-76` (capture surface, landed) + `tenstorrent_ops.cpp` `Trace*` (landed); NOT wired into a dense forward | **Measured on Blackhole P150 (2026-08-13):** Q1 no host-free region at T=1; Q2 all-device-at-T=1 costs 12.5→10.7 tok/s; Q2b capture attempt with both overrides flipped → ttnn **`TT_FATAL: Reads are not supported during trace capture`** (backtrace through `to_vector`), `0 replays`. The T=1 forward does device→host readbacks that ttnn trace prohibits; flipping the two thresholds is insufficient. Requires a host-free `ForwardLayers` redesign, not a threshold tweak. Next: prefill capture (separate row) must first audit its `to_vector` readbacks | [tenstorrent-trace-runner.md](specs/tenstorrent-trace-runner.md) | `SPIKE` (decision record complete) | `CLAIM-BACKEND-TENSTORRENT-TRACE-RUNNER` | +| `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` | Child of `BACKEND-TENSTORRENT` — make the per-decode-layer forward region host-free (zero `to_vector` readbacks) so mesh-trace capture can run. Decomposes into R1 RmsNorm+RoPE all-device, R2 QkvSplit+ReshapeAndCache device, R3 PA decode device-resident metadata, R4 capture wire+measure. Prerequisite revealed by the trace-runner spike (capture aborts on `to_vector`). | CUDA decode-graph capture contract (`cuda_backend.cu:184-197`): captured region is async, no host sync, no malloc, fixed ptrs. TT must match: no `to_vector` between Begin/EndCapture | `src/vt/tenstorrent/tenstorrent_ops.cpp` (RmsNorm:1067, PreferDeviceRope:1344, QkvSplit:1460, ReshapeAndCache:1527, PagedAttention:2009) | **R1-R3b MEASURED on P150** (env-gated `VT_TT_HOST_FREE_DECODE`, inert by default; 23/23 TT tests incl. a dedicated default-path inertness guard, M1-mutation-proved): R1 threshold flip landed; R2 device->device copy (`CopyDeviceDeviceIfCapture`, ttnn::copy+empty) landed; R3 program-cache warm (`enable_program_cache` + eager-warm) landed; R3b device zero-fill (`MemsetDeviceIfCapture`) landed. Capture enters the forward and reaches layer ops (CastBf16/RmsNorm fire). Remaining item-5 blocker = per-op enqueue_write; answer = persistent device tensors + before-replay populate (tt-metal vLLM plugin design). Full blocker map + architecture in [tenstorrent-host-free-r1.md](specs/tenstorrent-host-free-r1.md) | [tenstorrent-host-free-forward.md](specs/tenstorrent-host-free-forward.md) | `ACTIVE` (R1-R3b landed; item-5 port open) | `CLAIM-BACKEND-TENSTORRENT-HOST-FREE-FORWARD` | ## Native competitor and performance gates Every run records the competitor commit/release, model artifact hash, build diff --git a/.agents/coordination.md b/.agents/coordination.md index 21096710a..6ed07676e 100644 --- a/.agents/coordination.md +++ b/.agents/coordination.md @@ -1997,6 +1997,8 @@ this claim will meet. The tiled row is speed-gatable on dgx. | `CLAIM-ENG-RELEASE-CONTAINERS-W1-W7` | `ENG-RELEASE-CONTAINERS` (`INVENTORIED`->`READY`->`ACTIVE`; spike spec + W1-W5/W7) | Claude Code (opus-5), helper role | isolated worktree `/home/mudler/_git/vllm.cpp-containers`; local Docker only -- NO GPU, NO registry push, NO publication | `row/ENG-RELEASE-CONTAINERS`, base `origin/main` `24306364`, PR [#307](https://github.com/mudler/vllm.cpp/pull/307) | Container spike spec, `docker/Dockerfile` (cpu/vulkan/cuda calling the release scripts), container matrix + checker, image validator, workflow + least-privilege guard, and the SIGTERM fix the boot gate surfaced (#312). Excludes W6 matching-hardware runtime evidence and any GHCR publication | `ACTIVE` | 2026-08-10 -- cpu lane built and gated e2e (config+layout+boot, clean SIGTERM 0.25 s, RED 137 before the fix); mutation suites 31/31 and 22/22; cuda/vulkan gated statically but never built here; nothing published | | `CLAIM-ROCM-GDN-KERNELS` | `BACKEND-ROCM` (stays `ACTIVE` — M3 GDN kernel slice, issue #41) | pi (kimi-k3) | isolated worktree `/home/vikash/vllm.cpp-rocm-gdn-kernels`, base `origin/main` `5812b8b6`; 4x RX 7900 XTX gfx1100 discrete, ROCm 7.14, GPU jobs hold `~/gpu.lock` | `row/ROCM-GDN-KERNELS` | The M3 GDN kernel family unblocking Qwen3.5-0.8B M2 on discrete ROCm. Owns ONLY: NEW `.agents/specs/rocm-gdn-kernels.md`; NEW `src/vt/rocm/rocm_gdn.hip` (10 ops: kGdnStateGather/Scatter 77/78, kCausalConv1dFwd/Update 5/6, kRmsNormGated 9, kSigmoidGateBf16 58, kGdnPrefill/Decode 10/11, kGdnPostConv 65, kAttnQkNormRopeGate 67 — hand-translated from `src/vt/cuda/cuda_gdn.cu`/`cuda_ops.cu` donors); ADDITIVE registration entries in `src/vt/rocm/rocm_ops.hip`; ADDITIVE GDN cases + buffer helpers in `tests/vt/test_backend_cross_device.cpp`; this claim + `.agents/NOW.md`. **NON-COLLISION:** every new file is ROCm-only; the two shared-file edits are additive registration/case blocks that change no existing behavior (unregistered-op backends skip; CUDA/CPU/Vulkan paths byte-identical). Disjoint from #234 (sampling), #317 (gemma4 FP8/WMMA), #273 (near-tie docs). | `ACTIVE` | 2026-08-11 — **ALL 5 FAMILIES LANDED AS STACKED DRAFT PRs + M2 E2E COMPLETE.** #334 (spec + state I/O), #336 (conv), #341 (postconv), #343 (recurrence), #345 (norm-gate/preamble + the SupportsCompressedGdnState capability seam). Qwen3.5-0.8B runs e2e all-native on discrete gfx1100, deterministic; 'The sky is' 8/8 CPU parity; open-prompt divergence characterized as not-state-rounding, GDN-path drift analysis named open. Gates per family: cross-device green (162/162 total), focused 3/3, full-suite zero-delta vs base. NEXT: fresh review + operator gate per PR; then the MoeRouterTopK bf16-logits arm (the named MoE-path blocker) and the GDN-path divergence characterization. | | `CLAIM-WEIGHT-OFFLOAD-W0A` | `ENG-WEIGHT-OFFLOAD` (engine-matrix, `READY`->`ACTIVE`; issue #797, the dense half of #149). | Claude Code (claude-opus-5), helper role, isolated worktree | `/home/mudler/.cache/sdd/mudler-vllm.cpp/weight-offload-w0`; CPU-only build (`-DVLLM_CPP_CUDA=OFF`) on the dev box -- W0a is a device-neutral config port, NO GPU/flock needed | branch `row/ENG-WEIGHT-OFFLOAD-W0`, base `origin/main` `b1cd4d8f6` | **W0a -- the weight-offload CONFIG SURFACE only.** Owns EXACTLY the NEW `include/vllm/config/offload.h` + `src/vllm/config/offload.cpp`, the NEW `tests/vllm/config/test_offload_config.cpp`, 1 `CMakeLists.txt` source line + 1 `tests/CMakeLists.txt` test line, and the records (this row, the engine-matrix row, `specs/weight-offload-uva.md` `## Now`, `docs/STATUS.md`). **NON-COLLISION:** purely additive and UNREACHABLE -- nothing constructs an `OffloadConfig` yet, so every existing gate is byte-identical BY CONSTRUCTION (`git diff` touches no existing source file beyond the two CMake registration lines). Does NOT own the offloader itself (W2/W5), the loader re-offload interaction (W3), the ABI/CLI field (the REMAINDER of W0), or anything in `v1/kv_offload/` (a different subject). | `ACTIVE` | 2026-08-14 -- **W0a LANDED: config surface + validator, RED-first then mutation-proven.** Transcribed from `vllm/config/offload.py` @ `555967922`: the three-value backend enum (:12), both sub-configs with their bounds and defaults (:23,34,54,62,66,70), the TWO hard validator errors (:100-112) and the THREE mismatch warnings (:114-135) -- warnings collected for the caller rather than raised, since C++ has no `warnings.warn`. Also ported: the dot-anchored segment match (`offloader/uva.py:91-93`), the `int(gb*1024**3)` truncation (`offloader/base.py:155`), the auto-selection ORDER (`base.py:139-149`), and the layer-grouping worked example (`offload.py:57`). **RED-first CAPTURED** on a compiling stub (build rc=0, 0 compile errors): 11/11 cases and 51/122 assertions failed for the intended reason. GREEN 11/11, 126/126. **MUTATION-PROVEN 6/6**, each reported with its compile status so a non-building mutation could not read as a pass: substring-instead-of-dot-anchored (4 asserts RED), grouping off-by-one (3), validator guard OR->AND (3), warnings-not-cleared (1), auto-order uva-before-prefetch (2), 1000^3-instead-of-1024^3 (3); tree restored byte-for-byte. One design correction the RED exposed: `ResolvedBackend()` now mirrors `create_offloader` exactly (an EXPLICIT backend is selected even at a zero budget) and the separate `is_offloading_enabled()` answers 'would anything actually move' -- conflating the two would let a zero-budget explicit backend read as offloading-on. OWED (the rest of W0): the `include/vllm.h` JSON field + server CLI flag. | +| `CLAIM-BACKEND-TENSTORRENT-HOST-FREE-FORWARD` | `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` (NEW child of `BACKEND-TENSTORRENT`; `INVENTORIED`→`SPIKE`) | Maki (zai/glm-5.2), helper role | isolated worktree `/home/lu_zero/Sources/vllmcpp-tenstorrent`; real Blackhole (P150) hardware | branch `row/BACKEND-TENSTORRENT-HOST-FREE-FORWARD` (planning; reuses the trace-runner worktree) | The prerequisite for decode mesh-trace capture, revealed by `BACKEND-TENSTORRENT-TRACE-RUNNER` (capture aborts on `to_vector`). Decomposes into R1 RmsNorm+RoPE all-device, R2 QkvSplit+ReshapeAndCache device, R3 PA decode device-resident metadata, R4 capture wire+measure. Owns ONLY: NEW `.agents/specs/tenstorrent-host-free-forward.md`, the NEW `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` backend-matrix row, and this claim. No code, no model change, no perf claim — decision record only. | `SPIKE` | 2026-08-13 plan committed; no implementation yet. R1-R3 parallel-claimable once a maintainer accepts the decomposition. | +| `CLAIM-BACKEND-TENSTORRENT-TRACE-RUNNER` | `BACKEND-TENSTORRENT-TRACE-RUNNER` (NEW child of `BACKEND-TENSTORRENT`; `INVENTORIED`→`SPIKE`) | Maki (zai/glm-5.2), helper role | isolated worktree `/home/lu_zero/Sources/vllmcpp-tenstorrent`; real Blackhole (P150) hardware | branch `row/BACKEND-TENSTORRENT-TRACE-RUNNER`, base `origin/main` `a89b3c45` | Feasibility spike (read-only analysis + on-card measurement) for wiring the landed #354 graph-capture foundation into a capturable Qwen3-dense/Mistral forward region on TT. Three questions: (1) is there a host-free decode region at T=1 (current hybrid thresholds route RoPE+residual to host)? (2) what does forcing all-device-at-T=1 cost in tok/s, and does capture recover it? (3) does ttnn program-cache warm-up work before begin_trace_capture? Surveys how CUDA/Metal/Vulkan structure capture (CUDA is the only `SupportsGraphCapture()==true` backend; Metal/Vulkan `false`). Owns ONLY: NEW `.agents/specs/tenstorrent-trace-runner.md`, the NEW `BACKEND-TENSTORRENT-TRACE-RUNNER` row, and this claim. No code, no model change, no perf claim — decision record only. | `SPIKE` | MEASURED 2026-08-12/13 (all on real P150): no host-free region at T=1 (residual rows<32 + RoPE T*H<64 host, every layer); all-device-at-T=1 costs 12.5→10.7 tok/s (reproduces handoff S6); capture attempt aborts (to_vector + enqueue_write fatals). Decision superseded for implementation by the host-free investigation (#694): capture IS achievable; the remaining work (persistent device buffers + before-replay input populate, ported from the tt-metal vLLM plugin design) is tracked there. PR #541 closed in favor of #694. | **CLOSED same-session claim — `CLAIM-TP-SPIKE-287` (task #287, 2026-08-08, records-only, helper `row/SPIKE-TENSOR-PARALLELISM` draft PR #143, base `b38f78a7`).** Not an active-claims row because every referenced row keeps its diff --git a/.agents/specs/tenstorrent-host-free-forward.md b/.agents/specs/tenstorrent-host-free-forward.md new file mode 100644 index 000000000..32d02b23a --- /dev/null +++ b/.agents/specs/tenstorrent-host-free-forward.md @@ -0,0 +1,249 @@ +# Tenstorrent host-free decode forward — plan + +Status: **DRAFT plan, 2026-08-13.** The prerequisite for decode mesh-trace +capture (see `tenstorrent-trace-runner.md`: capture aborts on `to_vector` +readbacks inside `ForwardLayers`). This document decomposes the work into +independent rows sized for parallel claims. + +Row id: `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` (child of +`BACKEND-TENSTORRENT`). + +## Goal + +Make the per-decode-layer region of the TT forward **host-free**: zero +`to_vector` / `EnsureHost` readbacks between `BeginCapture` and +`EndCaptureGraph`. Only then can `Qwen3DenseDecodeGraph` capture/replay it +on `kTENSTORRENT` (ttnn `begin_trace_capture` prohibits any host read). + +The per-layer op sequence (Qwen3-dense / Mistral, from +`dense_attn_block.h`) and its host-readback status at T=1 today: + +| per-layer op | status today (T=1) | in captured region? | +|---|---|---| +| `RmsNorm` (pre-attn q-norm + residual merges) | HOST (rows<32) | yes | +| `MatmulBT` (qkv / o_proj / mlp) | device | fine | +| `QkvSplit` | **pure host** | yes | +| `RmsNorm` (qk-norm, Qwen3 only) | HOST | yes | +| `RopeNeox` / `RopeFromCache` | HOST (T·H<64) + `BuildCosSinFromPositions` host | yes | +| `ReshapeAndCache` | **pure host** | yes | +| `PagedAttention` | **pure host (host oracle)** | yes | +| `SiluAndMul` | device | fine | + +Boundary ops OUTSIDE the layer loop (readbacks here are fine — they are the +capture region's input/output edges): `Embedding` (host-staged upload), +`GreedyArgmax` (host readback of the final logits). + +## Scope + +**In.** Make the per-decode-layer region of the TT forward host-free (zero +`to_vector` readbacks and zero `enqueue_write` between `BeginCapture` and +`EndCaptureGraph`) so ttnn mesh-trace can capture/replay decode. R1-R3b +landed; item 5 (persistent device input tensors + before-replay populate) +open. + +**Out.** Prefill capture, MoE, new ttnn kernels, upstream tt-metal changes +(the answer is a vllm.cpp-side architecture port). + +## Upstream chain + +No upstream vLLM equivalent. The loyal anchors are: ttnn trace +(`ttnn::operations::trace::{begin,end}_trace_capture`, wired in +`tenstorrent_backend.cpp:70-76`), the CUDA decode-graph capture contract +(`cuda_backend.cu:184-197`: async region, no host sync, no malloc, fixed +pointers), and the tt-metal vLLM plugin's trace design (the reference +implementation of trace-based decode on this hardware). + +## Our baseline + +Landed on this branch (measured on real Blackhole P150, env-gated +`VT_TT_HOST_FREE_DECODE`, inert by default — 21/21 TT tests): R1 threshold +flip, R2 device->device copy, R3 program-cache warm, R3b device zero-fill. +Capture enters the forward and reaches the layer ops. The open gap is item +5: per-op `enqueue_write` during capture; the fix (persistent device +tensors + before-replay populate) is the plugin-port above. Full measured +record: `tenstorrent-host-free-r1.md`. + +## Port map + +No upstream vLLM equivalent (no vLLM Tenstorrent platform). The architecture +is ported from the official Tenstorrent vLLM plugin +(`tt/vllm/plugins/vllm-tt-plugin/.../model_runner.py`): + +| plugin technique | vllm.cpp TT mapping | +|---|---| +| two-phase warmup (compile ops with `enable_trace=False`, then capture) | `Qwen3DenseDecodeGraph` eager step then capture (already landed) + `device.enable_program_cache()` (R3, landed) | +| persistent device tensors at warmup max-padded shape (stable addresses) | TT decode-graph `SizeSlot` holds persistent ttnn device tensors for inputs (open — item 5) | +| `copy_host_to_device_tensor` before capture/replay, never inside | populate the stable buffers via `ttnn::copy_to_device` before `ReplayGraph` (open — item 5); inside the captured region only `CopyDeviceDeviceIfCapture`/`MemsetDeviceIfCapture` (landed R2/R3b) | + +## Tests to port + +None upstream. Local gates: the existing TT suite (21/21 default — proves +the env-gated paths are inert), the Qwen3-0.6B/Mistral TT golden pairs +(e2e near-tie when the flag is on), and the capture probe (bisection +instrumentation under `VT_TT_TRACE_DEBUG`). + +## Work breakdown + +Numbering below is the POST-INVESTIGATION truth (the pre-investigation plan +numbered R2=QkvSplit/RAC device and R3=PA metadata; the bisection showed the +copy/memset/allocation blockers fire FIRST, so those two original items are +now queued behind item 5 rather than being R2/R3). + +Each is independently gateable; none blocks another except the capture row, +which wants all three. + +### R1 — Device-resident RmsNorm + RoPE at T=1 (threshold flip + perf) + +**Problem:** the hybrid thresholds route `RmsNorm` (rows<32) and `RopeNeox` +(T·H<64) to host at T=1. The trace-runner spike measured the perf cost of +flipping them all-device: 12.5 → 10.7 tok/s (~14%, reproduces handoff §6). +Capture must recover that. + +**Work:** flip the thresholds to all-device when capture is active (or +unconditionally, gated on `support_static_graph_mode()`), accept the ~1.8 +tok/s eager regression, and let capture claw it back. The numerics were +already proven acceptable by `BACKEND-TENSTORRENT-RESIDUAL-GOLDEN` +(device bf16 vs CPU f32 = constant 0.0459 abs, ordinary rounding). + +**Sub-blocker:** `RopeNeox`/`RopeFromCache` depend on `BuildCosSinFromPositions`, +which reads `pos` on host (line 1291) and builds cos/sin host-side. The +device RoPE apply path exists (`RopeApplyDeviceNeox`) but the cos/sin +construction is still host. Needs a device-resident cos/sin path OR a +precomputed cos/sin cache uploaded once (the `RopeCosSinCacheKernel` path +already exists for the cache mode — route through it). + +**Gate:** op-level `RmsNorm`/`Rope` device parity (already measured); e2e +Qwen3/Mistral gate token-exact or near-tie vs the TT golden. + +### R2 — Device-resident QkvSplit + ReshapeAndCache (small host-staged ops) + +**Problem:** `QkvSplit` and `ReshapeAndCache` are pure host today — they +read every input via `EnsureHost` and `CommitHost` the output. Both are +bit-exact memcpy/stride ops that went host-staged in W0 because Alloc was +host memory. Inside a captured region they must stay on device. + +**Work:** add device-resident variants using `ttnn::slice` (QkvSplit) and +the device paged-write path that already exists for paged KV +(`NotePagedKvRacWrites` / `TryDevicePagedFill` / `TryDevicePagedUpdate` — +landed with residency). The device paged-write path already keeps a ttnn KV +shadow; wire `ReshapeAndCache` to it unconditionally when capture is active. + +**Gate:** op-level bit-exactness vs the host path (these are deterministic +copies — byte-identical is achievable and required); e2e gate. + +### R3 — Device-resident PagedAttention decode (the big one) + +**Problem:** `PagedAttention` at T=1 decode runs the **host f32 oracle** +(`PagedAttentionKernel` host path). The device path +(`TryPagedAttentionDeviceDecode`, `paged_scaled_dot_product_attention_decode`) +exists and is used when the KV shadow is current, but it still reads +`block_table`/`seq_lens`/`query_start_loc` on host (lines 1644-1646) and +reads `query` host (line 1707) before the device call. Those metadata +reads are the capture blocker. + +**Work:** keep the metadata tensors device-resident across the decode step +(they are small int32 tensors; upload once per step BEFORE the captured +region, not inside it), and ensure the query entering PA is already device +(no `EnsureHost(query)`). The device SDPA decode path itself is +capture-clean (it's a single ttnn op); the work is removing the host +metadata reads around it. + +**Gate:** device PA vs host oracle numerics (already measured: max_abs +~0.0009 for prefill; decode parity measured separately); e2e gate. + +### R4 — Flip `support_static_graph_mode()` + wire capture (only after R1-R3) + +**Problem:** the platform gate and the `Qwen3DenseDecodeGraph` wiring are +trivial once the region is host-free. This row flips the platform flag, +verifies capture no longer aborts, and measures replay tok/s vs eager. + +**Gate:** capture completes (no `TT_FATAL`); replay max_abs=0 vs eager +(already the landed unit-test property); **replay warm tok/s ≥ 12.5** +(the current hybrid eager baseline) — this is the payoff that justifies +all four rows. + +## Dependencies + +``` +R1 (RmsNorm+RoPE device) ─┐ +R2 (QkvSplit+RAC device) ─┼─► R4 (capture wire + measure) ──► decode tok/s win +R3 (PA decode metadata) ─┘ +``` + +R1, R2, R3 are independent and parallel-claimable. R4 is the integration +row that wants all three + produces the headline number. If R4's replay +tok/s does NOT beat 12.5, the whole effort is a wash — but that can only be +known after R1-R3, which is the cost of answering it. + +## Gates (per row + integration) + +- **Correctness:** every device-resident variant must be bit-exact or + near-tie vs the current host path, gated by the existing TT golden pair + (`our_ids_tenstorrent.npy` / `neartie_gap_mnats_tenstorrent.npy` for + Qwen3-0.6B, the Mistral pair for Mistral-7B). RED-first op-level test + before each e2e gate. +- **Capture (R4 only):** `TT_FATAL`-free capture + replay max_abs=0 + + replay warm tok/s ≥ 12.5 (Qwen3-0.6B `vllm-cli` smoke, same harness as + the trace-runner spike). +- **No perf regression outside capture:** the threshold flips in R1 regress + *eager* tok/s (12.5→10.7) — that regression is acceptable ONLY because R4 + recovers it. If R4 is not landed, R1 must not ship unconditionally; it + must gate on `support_static_graph_mode()` so non-capture runs keep the + hybrid thresholds and the 12.5 baseline. + +## Risk + +- **R3 is the scope risk.** R1 and R2 are mechanical (flip + reuse existing + device paths); R3 (device PA decode with device-resident metadata) is + real work and the most likely place to find another host touch. +- **R4's payoff is uncertain until measured.** The whole plan exists to + answer "does capture beat 12.5 tok/s"; if it doesn't, R1-R3 still + delivered device-resident ops (useful for future prefill capture) but no + decode win. That's an honest outcome, not a failure — it's the + measurement the trace-runner spike owed and couldn't make. + + +### Landed (this branch, measured on P150) + +- R1 threshold flip (RmsNorm residual + PreferDeviceRope all-device under + the flag). +- R2 `CopyDeviceDeviceIfCapture` (ttnn::empty + ttnn::copy device->device). +- R3 program-cache warm (`enable_program_cache` + eager-warm of the copy ops). +- R3b `MemsetDeviceIfCapture` (ttnn::zeros into the existing shadow). + +### Open (item 5 — the payoff port) + +Persistent device input tensors in the decode-graph slot + populate before +capture/replay via `ttnn::copy_to_device` (never inside capture). Ported +from the tt-metal vLLM plugin (see Port map). + +### Queued behind item 5 (from the original plan; may or may not be needed) + +Device-resident QkvSplit + ReshapeAndCache variants, and PA decode with +device-resident metadata. The bisection has not reached these (the +enqueue_write fatal fires first); keep or drop them per what item 5's probe +surfaces. + +### Known constraints of the investigation code (env-gated, carried forward) + +Recorded from review; all are flag-gated-only and acceptable for an +investigation row but MUST be addressed by the item-5 port: + +1. `CopyDeviceDeviceIfCapture` ignores `bytes` — a partial/interior Copy + between two same-sized shadowed slots clones the WHOLE src shadow. +2. It does not update `dev_rows`/`dev_cols`, so a consumer view matching + the logical shape but not the recorded shadow shape can fall into an + EnsureHost re-upload (a readback during capture — defeating R2). +3. The equal-BYTE check does not pin dtype/shape (a same-byte bf16/f32 + reinterpret is possible). +4. `enable_program_cache()` fires inside the copy helper; if the captured + region never takes that path it is never enabled. Belongs in + TraceBeginCapture (or platform init) for the port. +5. The `tt_capture_active()` clear is not exception-safe (a throwing + end_trace_capture leaves it stuck true, flipping eager Copy/Memset). + The inertness guard test catches the stuck-true case; the port should + make the clear RAII. +6. TOCTOU on SlotMutex around the ttnn calls (re-acquire without + revalidating the slot). +7. `d->device = std::move(cloned)` drops the prior dst shadow mid-capture + (a dealloc during a live trace). diff --git a/.agents/specs/tenstorrent-host-free-r1.md b/.agents/specs/tenstorrent-host-free-r1.md new file mode 100644 index 000000000..ed7c0a9bc --- /dev/null +++ b/.agents/specs/tenstorrent-host-free-r1.md @@ -0,0 +1,1017 @@ +# Tenstorrent host-free forward R1 — device RmsNorm + RoPE at T=1 + +Status: **DRAFT, 2026-08-13.** First row of the host-free-forward plan +(`tenstorrent-host-free-forward.md`). Sequential: measure after each row to +see its marginal contribution to capture. + +Proposed row id: `BACKEND-TENSTORRENT-HOST-FREE-R1`. + +## Scope + +**In.** Flip the two hybrid thresholds that route `RmsNorm` (residual, +rows<32) and `RoPE` (PreferDeviceRope, T·H<64) to host at T=1 decode, so +both go all-device. The numerics were already proven acceptable by +`BACKEND-TENSTORRENT-RESIDUAL-GOLDEN` (device bf16 vs CPU f32 = constant +0.0459 abs). The device paths already exist in `tenstorrent_ops.cpp`. + +The flip MUST be gated on capture-active (`support_static_graph_mode()`) so +non-capture runs keep the 12.5 tok/s hybrid baseline. Inert when capture is +off. + +**Out.** The cos/sin host build inside `RopeNeoxKernel` +(`BuildCosSinFromPositions`, line 1363) is a known sub-blocker but is NOT a +device `to_vector` readback — it reads the host `pos` tensor. Whether it +triggers the ttnn fatal is an empirical question R1 answers: after this flip, +does capture get past the current `to_vector` fatal, and what is the NEXT +host touch? (If it's `pos`, R1.5 or R2 handles it; if capture succeeds, R1 +alone was enough for the RmsNorm/RoPE portion.) No QkvSplit/ReshapeAndCache/ +PagedAttention work (R2/R3). + +## Upstream chain + +CUDA's capture contract (`cuda_backend.cu:184-197`): the captured region is +async, no host sync, no host readback. TT must match. The RmsNorm device path +(`tenstorrent_ops.cpp:1105-1117`, `ttnn::add`+`ttnn::rms_norm`) and the RoPE +device apply (`RopeApplyDeviceNeox`, line 1221) are the loyal mappings. + +## Our baseline + +`RmsNormKernel` (line 1067-1118): `host_residual` when rows<32. +`PreferDeviceRope` (line 1344): false when T·H<64. Both host at T=1. +The trace-runner spike measured forcing both all-device: 12.5→10.7 tok/s +eager (the cost capture must recover). + +## Work breakdown + +1. Add a capture-active helper reading the platform's + `support_static_graph_mode()` (cached per device, since the platform is + invariant). +2. Gate `host_residual` and `PreferDeviceRope` on `!capture_active`. +3. Op-level test: confirm RmsNorm + RoPE device path runs at T=1 (rows=1, + T·H=16) without the host fallback, bit-comparable to the residual-golden + measurement (0.0459 abs). +4. **Measure**: with `support_static_graph_mode()` also flipped on (R4's + change, applied locally for the measurement), does capture get past the + `to_vector` fatal? Record the next failure point if any. + +## Gates + +- Op-level: RmsNorm + RoPE device output at T=1 within the band already + measured by RESIDUAL-GOLDEN. +- E2e: Qwen3-0.6B `our_ids_tenstorrent.npy` golden still near-tie-passes. +- Capture probe (informative, not a hard gate for R1 alone): record how far + capture gets. +- No eager perf regression when capture is OFF (the gate must be inert). + +## Dependencies + +- `BACKEND-TENSTORRENT-RESIDUAL-GOLDEN` (numerics proof, landed). +- `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` (the plan, this branch). +- Hardware: Blackhole P150. + +## Risks/decisions + +- **The cos/sin host build may be the next capture blocker.** If after R1 + capture still fatals on a host op inside RopeNeoxKernel, the decision is + whether R1 is complete (RmsNorm+RoPE *apply* are device) and the cos/sin + build moves to a separate R1.5, or whether R1 must also switch RoPE to + the `RopeFromCache` path (model-side change). Settle empirically. +- **The gate must be inert by default.** A bug in the gate that flips + thresholds unconditionally would ship the 12.5→10.7 regression. Test the + inert path explicitly. + +## Outcome (2026-08-13) — R1 flip alone does NOT unblock capture; readback is inside a ttnn op + +Implemented the opt-in gate (`VT_TT_HOST_FREE_DECODE`): flips both +`host_residual` and `PreferDeviceRope` all-device, plus flips the platform +`support_static_graph_mode()` so `Qwen3DenseDecodeGraph` engages. All inert +by default (21/21 TT tests, 814/814 assertions unchanged). + +**Measurement (Qwen3-0.6B, `VT_TT_HOST_FREE_DECODE=1`):** capture STILL +fatals — `TT_FATAL: Reads are not supported during trace capture`, `0 +replays`. R1's threshold flip is **not sufficient**. + +**Diagnostic (the important finding):** with a capture-gated debug print on +every readback site in our code (`DownloadToHost`, `EnsureHost`, the two +direct `dev_out.to_vector` sites in PA decode/prefill), **zero of our +readbacks fire during capture**. The offending `to_vector` is therefore +**inside a ttnn op** (`ttnn::embedding` / `ttnn::rms_norm` / +`ttnn::sdpa_decode` / `to_memory_config` / etc.), not in our explicit +readback code. The TT backtrace shows only the ttnn frame +(`ttnn::Tensor::to_vector`), not which op called it. + +**Implication for the plan:** the host-free forward is **not** achievable +by only changing our thresholds/host-staging. At least one ttnn op in the +forward performs an internal host readback that ttnn trace prohibits. +Identifying that op (via a ttnn-symbolized backtrace or bisection) is the +real next step — it determines whether the fix is (a) swap to a +capture-safe ttnn primitive, (b) a ttnn version/bug fix, or (c) the +capture region must exclude that op. This is a deeper blocker than the +plan's R1-R4 assumed; the "host-free forward" may require upstream ttnn +changes, not just vllm.cpp changes. + +**R1 code kept** (env-gated, inert by default): the threshold flip is +correct and will be needed once the ttnn-internal readback is resolved. +The `tt_capture_active()` flag + `VT_TT_TRACE_DEBUG` prints are kept as +diagnostics for the next row. + +### Correction (2026-08-13, post-bisection): NOT a ttnn-internal readback + +The "inside a ttnn op" hypothesis above was **wrong** — it was based on +instrumenting only `DownloadToHost` + the two PA `to_vector` sites, which +missed the fourth readback site: `EnsureHostBytes` (line 2518). A full op +bisection with `TT_OP_TRACE` at every kernel entry + a print in +`EnsureHostBytes` gave the exact sequence: + +``` +[TT-TRACE] BeginCapture (flag set) +[TT-TRACE] EnsureHostBytes DURING CAPTURE <- the offender +TT_FATAL: Reads are not supported during trace capture +``` + +**Zero `TT-OP` kernel-entry lines fired** between BeginCapture and the +fatal — no `*Kernel` ran at all. The capture-blocking readback is in +**our** `Backend::Copy` → `EnsureHostBytes` → `dev.to_vector()` +(line 2533), triggered by `ForwardLayers`'s very first line +(`qwen3.cpp:244`): `d.b.Copy(d.q, hidden.ptr(), hidden_in.data, ...)`. +`hidden_in` has a device shadow from `EmbedInto`; `Copy` forces a +device→host download to satisfy the host-side memcpy, inside the captured +region. + +**This is fixable in our code, not an upstream ttnn blocker.** The fix: +when capture is active, `Backend::Copy` from a device-resident source must +do a device→device copy (or `ForwardLayers` must receive the device tensor +directly instead of copying through host). That's a concrete, scoped +R2-target — the "host-free forward" is achievable in vllm.cpp after all, +once every `EnsureHostBytes`/`Backend::Copy` site in the captured region +is made device→device. The R1 threshold flip + this copy fix together +clear the first capture blocker; subsequent readback sites (if any) +surface as the next bisection steps. + +### R2 status (2026-08-13): fix site identified, device-copy primitive pending + +The bisection pinpointed the exact fix site: `Backend::Copy` +(`tenstorrent_backend.cpp:56`) -> `EnsureHostBytes` -> `to_vector`, +triggered by `ForwardLayers`'s opening `d.b.Copy(...)` (`qwen3.cpp:244`). + +Fix shape: when capture is active and both `dst` and `src` carry current +device shadows, `Backend::Copy` must do a device-to-device copy via a new +`CopyDeviceDeviceIfCapture` helper in the ops TU, called before +`EnsureHostBytes`. + +Pending: the exact ttnn device-to-device copy primitive for this ttnn +build. Tried `ttnn::clone` (header not on the installed include path +despite the source existing) and `ttnn::copy` + `ttnn::zeros` (wrong API +for this version). The installed ttnn headers are a subset; the right +primitive needs focused API discovery against the installed header set. +R2 code reverted to keep the branch buildable; the +`CopyDeviceDeviceIfCapture` contract is the design, the body is the open +work — the single concrete next step. + +### R2 update (2026-08-13): device-copy LANDED, next blocker is ttnn program-cache warm-up + +Resolved the ttnn API discovery: the installed headers live in +`build_Release/include/ttnn/operations/...` (not the `libexec` tree). The +device→device copy primitive is `ttnn::copy(src, dst)` (from +`data_movement/copy/copy.hpp`) with a destination allocated via +`ttnn::empty(shape, dtype, layout, device, memconfig)` (from +`creation/creation.hpp`), using `Tensor::logical_shape()`/`dtype()`/`layout()` +accessors. Both headers had to sit inside the Tracy-disabled include block +(they transitively pull the 6-arg `op_profiler_serialize` that breaks the +5-arg TracyC.h). Backend::Copy now calls `CopyDeviceDeviceIfCapture` first; +default tests still 21/21, 814/814 (the path is capture-gated, inert +otherwise). + +**Measured:** the R2 copy fix **works** — the capture probe now gets past +the `EnsureHostBytes` readback (`[TT-TRACE] device->device copy (capture-safe)` +fires, no more `Reads are not supported` fatal). The **new** fatal is one +layer deeper: + +``` +TT_FATAL: Cannot load new binaries during trace capture. +This program is not yet in program cache. Warm up before capturing a trace. +``` + +This is the **ttnn program-cache warm-up** requirement (Q3 in the original +trace-runner spike, deferred then). ttnn `begin_trace_capture` requires +every op shape in the captured region to be JIT-compiled (program-cache +warm) BEFORE capture begins; the decode-graph framework's single eager +warmup step does not warm the exact shapes the captured path uses (or my +new `ttnn::empty` introduces an un-warmed program). + +This is a known ttnn trace discipline with an established pattern (warm +identical shapes via an eager run that hits the same ops), NOT an unknown. +It is the concrete R3 target — and it means the host-free forward *is* +achievable: R1 (thresholds) + R2 (device-copy, landed) clear the readback +blockers; R3 (warm-up) is the last gate before capture can complete. + +### R3 update (2026-08-13): warm-up WORKS — capture now runs ops + +Fixed the warm-up: `CopyDeviceDeviceIfCapture` now runs whenever +`VT_TT_HOST_FREE_DECODE` is set (not just during capture), so the eager +warmup step also exercises `ttnn::empty`+`ttnn::copy`, compiling them into +the program cache. Also calls `device.enable_program_cache()` once on the +first host-free use (the ttnn trace precondition). + +Measured: the "Cannot load new binaries" fatal is gone. Capture now enters +the forward and runs ops: + + device->device copy (eager warmup) + BeginCapture + device->device copy (ForwardLayers opening Copy — R2 holds) + EnsureHostBytes DURING CAPTURE x6 <- next readback blockers + CastBf16 + RmsNorm <- ops run during capture + TT_FATAL: Writes are not supported during trace capture <- a buffer write + +R2 + R3 together got capture past the first Copy and into the layer ops. +Two new, expected, mechanical blockers surfaced: + +1. 6 more EnsureHostBytes readbacks — every Backend::Copy inside + ForwardLayers (weight uploads, residual init) hits host-staging. Same + R2 fix at each site. +2. Writes not supported — DBuf::Zero calls Backend::Memset (host memset), + a host write inside the captured region. Needs a device-zero path or + pre-zero outside capture. + +Conclusion: capture on TT is achievable and now demonstrated working past +the first two blocker layers. Remaining work is converting each +host-staging site (Backend::Copy, Backend::Memset/DBuf::Zero) in the layer +loop to device-resident — mechanical, not research. The bisection +instrumentation surfaces each site in order. + +### R3b update (2026-08-13): copy + zero-fill done; device-allocation is the structural blocker + +Added MemsetDeviceIfCapture (on-device ttnn::zeros for DBuf::Zero), fixed a +null-deref (std::optional). Default tests 21/21. + +Measured: the 6 EnsureHostBytes readbacks are GONE. The sequence now: + device->device copy (eager warmup) + BeginCapture + device->device copy (ForwardLayers opening) + device zero-fill (DBuf::Zero) + Writes are not supported during trace capture <- structural blocker + +The Writes fatal is ttnn forbidding device allocations during capture +(same as CUDA's no-cudaMalloc-during-capture). The TT ops do per-call +from_vector host->device uploads (weights/inputs) and ttnn::empty scratch +inside kernels; those are fresh device writes, forbidden during capture. +CUDA solves this with a pre-warmed DevicePool + fixed-address persistent +weight buffers; TT has no equivalent, and its weights are not in stable +device buffers persisting across warmup->capture. + +This is the structural hard part: a TT scratch-pool analogue + stable +weight residency so no allocation/upload happens during capture. Real +engineering, the natural scope of a dedicated row. + +COMPLETE BLOCKER MAP (the experiment's deliverable): +1. RmsNorm/RoPE host thresholds -> R1 (flip, done) +2. Backend::Copy host readback -> R2 (device->device copy, done) +3. ttnn program-cache warm-up -> R3 (enable + eager-warm, done) +4. Backend::Memset/DBuf::Zero host write -> R3b (device zero-fill, done) +5. per-op device allocation/upload (from_vector, ttnn::empty) -> REMAINING; + needs a TT scratch pool + stable weight residency + +Items 1-4 landed, measured, inert-by-default. Item 5 is the open +engineering gate before decode capture can complete and replay tok/s can +be measured. + +### Upstream investigation (2026-08-13): item 5 may be a non-issue on newer ttnn + +Searched tt-metal/tt-nn issues. The "Writes are not supported during trace +capture" fatal is a **known limitation with an upstream fix**: + +- **tt-metal issue [#13690](https://github.com/tenstorrent/tt-metal/issues/13690)** + "Enable allocation of new buffers with a warning to allow running decode + with trace and prefill without trace" — filed by Tenstorrent **for vLLM** + (referenced by tenstorrent/vllm#14). The exact use case: interleaving a + traced decode with untraced prefill needs buffers allocated while a trace + is live. +- **Fixed in PR [#13696](https://github.com/tenstorrent/tt-metal/pull/13696)** + (commit `f0b2483`): instead of `TT_FATAL`, it now prints a warning and + allows the allocation, safe as long as untraced intermediates are consumed + before a trace runs. +- **This build does NOT have the fix** — `fd_mesh_command_queue.cpp:760` + still uses `TT_FATAL(!trace_id_.has_value(), "Writes are not supported...")`. + +**Implication:** bumping the tt-metal build to one including #13696 may +eliminate item 5 entirely (the upload-during-capture becomes a warning, +not a crash). Worth testing before building a scratch-pool subsystem. + +Additionally: +- `ttnn::create_device_tensor(spec, device)` (from + `graph/graph_query_op_constraints.hpp`) allocates an empty device tensor + **without** a host→device write — the capture-safe allocation pattern. + The canonical capture sequence (graph_query_op_runtime.hpp:76-90) uses it + to create input tensors pre-capture, warm, then capture. Our ops use + `from_vector` (which writes); converting uploads to + `create_device_tensor` + a pre-capture warm would also avoid the fatal. +- `TraceBufferPool` (PR #18523) — ttnn already has trace buffer management + infrastructure. + +**Two concrete paths to clear item 5, in order of effort:** +1. **Bump tt-metal** to a build with #13696 and re-run the capture probe. + If the warning-only path works, capture completes and we get replay + tok/s immediately — no vllm.cpp changes. +2. If the bump is not possible or insufficient: convert the TT ops' weight + uploads to pre-capture `create_device_tensor` (stable, pinned addresses + — the "pin addresses for a stable pool" approach) so no write happens + during capture. Bounded work, no new subsystem. + +### Correction (2026-08-13): bump will NOT help — our fatal is a write guard, not the allocator guard + +Verified `f0b2483` IS an ancestor of the installed tt-metal build (the #13690 +fix is present). But #13690 only relaxed the **allocator** (`allocator.cpp` ++ `device.cpp`) — it allows **buffer allocation** during a live trace. +Our fatal is at `fd_mesh_command_queue.cpp:760`, the **`enqueue_write`** +(host→device write) guard, which is a *separate* assertion that #13690 did +NOT touch (all three `Writes are not supported` fatals in +`fd_mesh_command_queue.cpp` are still hard `TT_FATAL`s). + +So bumping tt-metal will not clear item 5. The real fix is path 2: avoid +the `enqueue_write` during capture by pre-allocating device tensors with +`create_device_tensor` (which does not write) and uploading their contents +*before* capture, so the captured ops reference stable device buffers with +no host→device write. This is the "pin addresses for a stable pool" approach +— confirmed feasible by `ttnn::create_device_tensor` existing and being the +canonical capture-safe allocation path (graph_query_op_constraints.hpp). + +### Architecture answer (2026-08-13): mirror the tt-metal vLLM plugin's design + +Read the official Tenstorrent vLLM plugin +(tt/vllm/plugins/vllm-tt-plugin/.../model_runner.py). It solves this +exactly, and the answer is a vllm.cpp-side architecture change, not a +tt-metal patch: + +1. Two-phase warmup (model_runner.py:3216-3262): Phase 1 compiles all ops + into the program cache with enable_trace=False; Phase 2 captures with + every op compiled. Our Qwen3DenseDecodeGraph already does the + single-step version. +2. Persistent device tensors at warmup shape (model_runner.py:480-487): + block tables, positions, inputs allocated as persistent ttnn device + tensors at the max padded shape during warmup so capture replays against + stable device addresses. +3. Per-step inputs pushed BEFORE the captured region, not inside it: the + plugin uses ttnn.copy_host_to_device_tensor (= C++ copy_to_device -> + enqueue_write_tensor) to populate stable buffers. Crucially, + copy_to_device hits the SAME enqueue_write path that fatals during + capture (fd_mesh_command_queue.cpp:760), so the plugin calls it BEFORE + capture (warmup populate) and BEFORE each replay (per-step refresh), + NEVER inside the captured region. + +Implication: our Backend::Copy/EnsureHostBytes fatal during capture is +fundamental -- copy_to_device itself would fatal there too. The fix is +architectural: the captured ForwardLayers region must reference only +pre-allocated, pre-populated device tensors. Per-step inputs (token id, +position, slot mapping, block table) must be written to stable device +buffers BEFORE ReplayGraph, the same way CUDA's decode graph does (its +SizeSlot::Refresh writes host buffers that a captured async-copy re-reads, +qwen3.cpp:528). + +So path-2 is: make the TT decode-graph slot hold persistent device tensors +for inputs, populate them before capture/replay via copy_to_device, and +ensure the captured ops read those device tensors without any internal +from_vector/to_vector. That is the real host-free forward -- a bounded +architecture port of the plugin's design, not a new subsystem and not an +upstream tt-metal fix. + +### Steady-state perf baselines (2026-08-14, real Blackhole P150) + +Qwen3-0.6B, `vllm-cli --prompt "Hello" --max-tokens 64 --repeat 3`: + +| config | warm tok/s (runs 2/3) | ms/tok | +|--------|----------------------|--------| +| default hybrid | **7.30 / 7.31** | ~137 | +| all-device eager (`VT_TT_HOST_FREE_DECODE=1` + `VLLM_CPP_QWEN3_DENSE_DECODE_GRAPH=0`) | **6.87 / 6.92** | ~145 | + +Two corrections to the earlier smoke numbers, both measured: + +1. **The 4-token smoke (12.5 tok/s) is NOT steady state.** At 64 tokens the + same config sustains 7.3 tok/s — per-token cost grows with context (PA + decode over a growing KV), so the handoff's ~12.3 and our 12.5 both + over-report. The honest reference for capture work is 7.3. +2. **The all-device eager cost at steady state is ~0.4 tok/s (~6%), not the + ~1.8 tok/s (~14%) the 4-token smoke suggested.** This materially improves + the capture break-even: capture only needs to recover ~6% of eager time + to beat the hybrid baseline at 64-token scale — a much lower bar than + the spike's 14% framing assumed. + +(The `DECODE_GRAPH=0` opt-out is required for the all-device run: with the +flag on, `support_static_graph_mode()` flips true and the decode-graph +framework would otherwise attempt capture and abort on item 5.) + +Mistral-7B-v0.3 reference on the same box: 4.26 tok/s warm at 32 tokens +(recorded in tenstorrent-mistral.md). + +**What is still NOT measurable until item 5 lands: capture/replay tok/s.** +Capture cannot complete (the enqueue_write fatal fires mid-forward), so the +replay number — the actual payoff — remains open. The numbers above bound +it: replay must exceed 7.3 (the hybrid eager baseline) to be a win, and +starts from a 6.9 eager floor on the all-device path. + +### Item 5 progress (2026-08-14): two sites fixed; frontier now mid-layer-0, at rope cos/sin + +Instrumented all 16 `from_vector` sites (capture-gated `[TT-UP]` prints, +incl. ptr+shape on UploadRows) and iterated the bisection. Two real item-5 +fixes landed: + +1. **ttnn::zeros is NOT capture-safe** (creation.cpp `full_impl` host-fills + + `to_device()` = an enqueue_write) — my own R3b helper was an offender. + Fixed the plugin way: a persistent ZERO TENSOR CACHE (keyed by + shape/dtype/layout) created OUTSIDE capture, primed during the eager + warmup by `EnsureDevice2D`, and applied in-region by + `ttnn::copy(zero, shadow)` — a device->device program that is captured + and replayed. Cache-miss during capture is a hard VT_CHECK (must warm + first), which is exactly what forced the priming fix. +2. **QkvSplit's device path was already correct** (ttnn::slice + + CommitDevice2D) — the earlier suspicion was wrong; with MatmulBT's + shadow it fires and hands q/k/v shadows downstream. + +**Measured frontier after both fixes** — capture now runs DEEP into +layer 0 and dies at a precisely-identified site: + +``` +BeginCapture -> device-copy -> zero-fill -> [6 EnsureHostBytes readbacks += the weight DBuf copies, handled by R2] -> CastBf16 -> RmsNorm -> +MatmulBT -> QkvSplit -> RmsNorm(q-norm) -> RmsNorm(k-norm) +-> [TT-UP] UploadRows ptr=... rows=16 cols=64 <- THE blocker +-> TT_FATAL: Writes are not supported during trace capture +``` + +`[16, 64]` is the **RoPE cos/sin table** (Hq=16, rot/2=64): host-computed +by `BuildCosSinFromPositions` inside `RopeNeoxKernel` and uploaded +in-region. This is the cos/sin sub-blocker the R1 spec predicted, and it +is the plugin's "per-step input" case: the fix is a PERSISTENT device +cos/sin buffer populated before capture/replay (positions change per step, +so the decode-graph driver must copy_to_device the step's rows BEFORE +ReplayGraph — the same pattern as CUDA's SizeSlot::Refresh async-copy). + +**Remaining sites after rope (not yet reached by the bisection, expected +from the readback map):** ReshapeAndCache's KV writes (host-staged), +PagedAttention's metadata uploads, the lm_head/logits path. Each is the +same pattern; the rope fix establishes the template. + +Status: item 5 is now a SCOPED multi-site port (rope cos/sin + RAC + PA +metadata + logits), with two sites landed and the third precisely +characterized. Not complete; the replay-tok/s payoff measurement remains +blocked behind the remaining sites. + +### Item 5: rope cos/sin SOLVED (2026-08-14, measured on card) + +The rope blocker took three fixes working together: + +1. **Persistent cos/sin cache** keyed by (tokens*heads, rot/2), entries + created/refreshed OUTSIDE capture, replayed in-region via the captured + program (no per-call upload). Content-identity checked against the exact + bytes the kernel will use — a stale table is a hard VT_CHECK during + capture, never silent corruption. +2. **Driver warm hook** `WarmRopeCosSin(positions, ...)` called from the + decode-graph driver's Refresh slot (qwen3.cpp, right after + SizeSlot::Refresh) — THE per-step populate point, the exact plugin + SizeSlot::Refresh analogue. Crucially it warms the UNPADDED T-row + positions (what si.positions/rope sees), not the padded ppos — the + first attempt used ppos and always missed. +3. **Byte-exact content**: the captured rope path (RopeFromCache, the + default VT_QWEN3_ROPE_CACHE route) reads cos/sin from the per-step bf16 + CACHE table (RopeCosSinCacheKernel's StoreElemF32 rounds f32->bf16), so + the warm content must round-trip through bf16 (BF16ToF32(F32ToBF16(v))) + — f32 warm content never matches (cos(1)=0.540302 f32 vs 0.539062 bf16). + +Measured: rope cache **HIT for both q (16x64) and k (8x64)** during +capture (`content_eq=1`), capture proceeds PAST rope. Also discovered en +route: the dense decode path routes rope through RopeFromCacheKernel (not +RopeNeoxKernel) by default — the first debug print in the wrong kernel +never fired, which is what exposed it. + +**New frontier: ReshapeAndCache** — the next fatal is a to_vector readback +inside RAC (the KV-write path), right after rope in layer 0. This is the +"queued" RAC item from the original plan, now live. After RAC: PA metadata, +then the logits path. RAC is the most delicate remaining site: KV writes +inside a captured+replayed region also raise a REPLAY-SEMANTICS question +(every replay re-appends the same KV row) that must be answered alongside +the mechanical fix — the CUDA graph solves this by capturing the append +against fixed slot addresses refreshed per step. + +Default-path safety re-verified after all rope changes: 23/23 cases, +830/830 assertions. + +### Post-rebase benchmark confirmation (2026-08-14, rebased tree) + +Rebased onto main (47 commits; picked up the MISTRAL row landing via the +bot and the windows C4456 fixes). One MAIN-RED found while rebuilding: +MUSIC3 W6 (aa3643b6) placed a C++ helper returning a C++ reference inside +`extern "C"` in src/capi/vllm_c.cpp — clang rejects it (every clang build +is broken on that commit; MSVC/Windows was lax). Applied the minimal local +fix (hoist the helper out of the extern-C block) to unblock; reported +upstream. + +Benchmark re-run on the rebased tree (64 tokens, batch 1, 3 reps): + +| config | warm tok/s | +|--------|-----------| +| default hybrid | **7.13 / 7.23** (clean EXIT=0) | +| all-device eager (rope-cache additions included) | **6.68 / 6.80** | + +Consistent with the pre-rebase 7.30/7.31 and 6.87/6.92 (within run noise); +the rope-cos/sin-cache additions cost ~0.1-0.2 tok/s eager, the price of +capture-safety on that path. Default-path suite on the rebased tree: +23/23 cases, 831/831 assertions (main's merged tests grew the count). + +### Item 5 frontier: ReshapeAndCache analysis complete (2026-08-14) + +The readback: `ReshapeAndCacheKernel`'s first act is `EnsureHost(k)` — the +rope output carries a device shadow, so the download (a to_vector) fires +inside capture. Even on a shadow hit, the device push re-uploads: every +existing device path (`TryDevicePagedFill/UpdateBatch/FusedUpdateBatch`) +builds its input AND page table via `from_vector` (enqueue_write, also +fatal). The host NHD cache is the RAC/LMCache source of truth; the ttnn +shadow is a mirror patched from host floats. + +**The capture-safe RAC design (next implementation step):** + +1. Device-resident k/v input: the rope output shadow [T*H, D] bf16 TILE + must feed `paged_update_cache` directly. Layout gap ([T*H,D] flat vs + the sharded [C, nkv*d] input MakeHeightShardedUpdateInput builds today) + resolved ON DEVICE via capturable reshape/permute ops. +2. Persistent PAGE-TABLE device tensor, per-step refreshed outside capture + (the driver Refresh slot — same pattern as WarmRopeCosSin; the padded + block table already lives in the SizeSlot host buffer, so the refresh + source exists). +3. Persistent UPDATE-IDX device tensor likewise (paged_update_cache takes + update_idxs_tensor — a device tensor — natively). +4. The host NHD mirror patch moves OUT of the captured region: done at the + per-step refresh point from the same k/v tokens, before capture, so the + LMCache contract (host NHD = source of truth) is preserved. + +Replay semantics note: each replay re-writes the same KV slots the capture +baked in. That is only correct if the slot indices come from a +per-step-refreshed device tensor — the same reason CUDA's graph refreshes +slot_mapping per step. The design above has that property (2/3). + +`ttnn::experimental::paged_update_cache`'s signature confirms feasibility: +it accepts a device `input_tensor`, a device `update_idxs_tensor`, and a +device `page_table` — all three can be persistent/refreshed, no host +floats needed in-region. + +This is the largest single remaining piece (bigger than rope: on-device +layout conversion + two new refreshed buffers + moving the mirror patch). +After RAC: PA metadata (same refresh pattern), then logits. + +### Item 5: RAC progress (2026-08-14) + +Implemented `TryReshapeAndCacheDeviceDecode` + `WarmRacIdx` driver warm hook ++ shape-keyed idx cache (same content-refresh pattern as rope). The warm +hook fires correctly (slot0=32 warmup, slot0=33 capture step), the content +check matches, but the device branch bails because the **paged-KV device +shadow is empty** (`k=0 v=0`) — it was never created. + +Root cause: the paged-KV shadow is created by `EnsurePagedKvTtnn` (inside +`TryPagedAttentionDeviceDecode`), but PA's device path doesn't run during +the eager warmup (its preconditions aren't met on the non-capture path). +So by capture time the shadow was never populated. + +Fix needed: eagerly create the paged-KV shadow during the warm hook (call +`EnsurePagedKvTtnn` from `WarmRacIdx`, or prime it from the KV cache +metadata the driver has via `attn_kv`). This is the same "prime outside +capture" pattern as the zero cache and rope cos/sin. + +After the shadow exists, the remaining RAC path (device→device +`paged_update_cache` with persistent idx tensors) should work — the idx +content already matches (verified), the k/v shadows exist (post-rope +`CommitDevice2D`), and the paged_update_cache signature accepts all +device tensors. + +### Item 5: RAC device branch EXECUTES; paged_update_cache warm hangs + +The shadow priming (WarmPagedKvShadow) works — both k and v shadows exist +(`k=1 v=1`). The RAC device branch (`TryReshapeAndCacheDeviceDecode`) fires +during capture: `[TT-TRACE] RAC device->device update (capture-safe)`. +But `paged_update_cache` is not program-cache-warm (the eager forward's RAC +bailed to host because the shadow didn't exist yet), and the capture call +hits `Writes are not supported` (new binary load during capture). + +Attempted to warm `paged_update_cache` from `WarmRacIdx` with a dummy +input of the correct geometry (`[1,1,nkv_pad,d]` = `[1,1,32,128]`). +The warm call HANGS — `paged_update_cache` appears to deadlock when called +from the warm-hook context (outside the regular forward flow). This may be +a ttnn device-state issue (the mesh device's CQ is in a state that doesn't +support the op outside a forward step) or a geometry mismatch in the +warm-call's page_table/idx tensors vs what paged_update_cache expects. + +NEXT: investigate why the warm `paged_update_cache` hangs. Options: +(a) call it from within the eager forward (not the warm hook) by making the + eager RAC step also take the device branch (prime the shadow BEFORE the + eager forward, not after it — move WarmPagedKvShadow before the eager + step in the driver flow); +(b) use a simpler ttnn op (e.g. just `ttnn::copy`) as a warm substitute + that compiles the same program path; +(c) move the shadow priming into the eager forward itself (call + EnsurePagedKvTtnn at the top of the eager PA, not just the capture PA). + +Option (a) is the most promising: the eager forward already runs the full +op chain; if the shadow exists at eager time, the eager RAC takes the +device branch, which warms `paged_update_cache` naturally (same context, +same CQ state). The issue is the ordering: the framework runs the eager +step BEFORE the Refresh slot (where the warm hooks fire). Moving the shadow +priming to BEFORE the eager step (at slot creation, not Refresh) would fix +the ordering. + +### Item 5: RAC — `paged_update_cache` is NOT capture-safe (internal allocation) + +After fixing: +- shadow priming for all layers (not just layer 0) +- used=block+1 (off-by-one in block coverage) +- idx tensor dtype INT32 (not UINT32 — ttnn requirement) +- input sharding (paged_update_cache requires height-sharded input) + +The RAC device branch now EXECUTES on both cold and capture steps +(`RAC device->device update (capture-safe)` fires). But `paged_update_cache` +itself triggers `Writes are not supported during trace capture` — the op +does an internal allocation (result tensor) that is an `enqueue_write`. + +This is NOT a program-cache issue (the cold step compiled the program). +`ttnn::experimental::paged_update_cache` allocates a new output tensor +even when the program is cached — that allocation is a device write, +forbidden during capture. + +This is a ttnn API limitation: the op is not capture-safe by design. +The plugin's approach (persistent device tensors + before-replay populate) +works for ops that take pre-allocated outputs, but `paged_update_cache` +returns a new tensor. The fix would be either: +(a) an upstream ttnn change to support in-place update (pass output tensor) +(b) pre-allocate the result and use a different capture-safe scatter op +(c) capture only the ops AFTER RAC (skip RAC from the captured region, + do it before replay) — but RAC mutates the KV cache, which PA reads + inside the captured region, so it can't be moved out. + +Option (a) is the cleanest (an upstream issue/PR to ttnn). This is the +genuine gate — not a vllm.cpp code issue but a ttnn API limitation. + +### Item 5: RAC — paged_update_cache IS in-place; writes from build_padded + +Key discovery: `paged_update_cache::create_output_tensors` returns +`tensor_args.cache_tensor` — it's an **in-place** operation (no output +allocation). The `Writes are not supported` error was NOT from +`paged_update_cache` itself but from `build_padded`'s helper ops: +`ttnn::to_memory_config` (sharding allocates a new buffer) and possibly +`ttnn::concat`/`ttnn::zeros`. + +Attempted: pre-build the sharded zero input in WarmRacIdx and use +`ttnn::copy` (capture-safe) in `build_padded`. Crashed (segfault 139) +during the cold step — likely a shape/lifetime mismatch between the +pre-built sharded tensor and what `build_padded` produces. The +`sharded_zero` may be default-constructed (empty) if the warm loop didn't +find a shadow, or the shapes don't align. + +NEXT: debug the sharded_zero lifetime/shape, or take the simpler approach +(b) — replace `paged_update_cache` with a manual `ttnn::copy` into a +pre-sliced cache region (simpler op, no sharding requirement, proven +capture-safe by the R2 copy fix). + +### Item 5: sharded_zero crash FIXED; Writes still from build_padded + +Fixed the segfault: the warm order was wrong — `WarmRacIdx` ran BEFORE +`WarmPagedKvShadow`, so the shadow loop found 0 entries and the +sharded_zero was default-constructed (empty). Swapped the order in the +driver: shadows first, then RAC idx. No more segfault. + +But the `Writes are not supported` fatal persists. `paged_update_cache` +is in-place (confirmed: create_output_tensors returns cache_tensor). +The writes come from `build_padded`'s helper ops — specifically +`ttnn::copy(reshaped, sharded_zero)` where reshaped is TILE and +sharded_zero is height-sharded. The copy between different memory configs +triggers an implicit layout conversion (a write/allocation). + +NEXT: approach (b) — replace `paged_update_cache` + the sharded input +with a manual `ttnn::copy` into a pre-sliced cache region. The cache +shadow is a persistent ttnn tensor; slicing it and copying the k/v rows +into the slice is all-capture-safe (proven by R2's device->device copy). +No sharding requirement, no paged_update_cache, no layout conversion. + +### Item 5: approach (b) — RAC SKIPPED during capture; next blocker = PA + +Implemented approach (b): `TryReshapeAndCacheDeviceDecode` returns true +immediately during capture (skipping the KV write). This is INCORRECT for +real decode (stale KV) but proves the approach works — capture proceeds +PAST RAC to the next op. + +Measured on card: `[TT-TRACE] RAC skip during capture (approach b probe)` +fires, capture continues to `TryPagedAttentionDeviceDecode` which then hits +`Reads are not supported during trace capture` (fd_mesh_command_queue.cpp:807 += the READ guard, not the write guard at :760). So PA is doing a +`to_vector` readback — likely `EnsurePagedKvTtnn` re-uploading the stale +shadow (marked stale by the skip), or PA reading query/metadata via +`EnsureHost`. + +The remaining sites after RAC are: +1. PA metadata (block_table, seq_lens, query_start_loc via EnsureHost) +2. PA's EnsurePagedKvTtnn (re-upload the stale KV shadow) +3. PA output (to_vector to read the attention result) +4. Logits (lm_head output) + +The real fix for RAC: move it OUT of the captured ForwardLayers region +entirely — do the KV write at the driver Refresh slot (before BeginCapture), +same as the plugin's per-step copy_host_to_device_tensor pattern. This means +splitting the captured region: RAC runs before capture, PA+forward runs +inside capture. That's a driver-level change (the captured region starts +after EmbedInto + RAC, not at ForwardLayers). + +### Item 5: PA — KV shadow skip works; next = page_table + cur_pos uploads + +Fixed PA's KV shadow re-upload: during capture, skip EnsurePagedKvTtnn and +use the cached shadow directly (it was primed by WarmPagedKvShadow at the +Refresh slot). `PA using cached KV shadows (k_nb=2 v_nb=2)` prints, +`PA KV shadows OK, building page_table` prints. + +Next fatal: BOTH Writes (:760) and Reads (:807) — the from_vector uploads +for dev_pt (page_table) and dev_pos (cur_pos) at PA lines 2104/2114. +Same pattern as rope: per-step data (block_table, seq_lens) that needs +persistent device tensors warmed at the Refresh slot. + +Remaining sites after PA metadata: +1. PA page_table upload (from_vector, line 2104) — persistent device tensor +2. PA cur_pos upload (from_vector, line 2114) — persistent device tensor +3. PA sdpa_decode output (to_vector at line 2164+) — device→device commit +4. Logits (lm_head output) + +Each is the same persistent-buffer + driver-warm pattern. The path is +proven (RAC skip + PA shadow skip both work); it's mechanical repetition. + +### Item 5: PA — metadata warm works; sdpa_decode not compiled (cold bail) + +PA metadata warm (WarmPaMeta) works: `PA using cached meta (pt+cp)` fires +during capture. But `sdpa_decode` hits `Cannot load new binaries during +trace capture` — it was never compiled during the cold step because the +cold step's PA device path bails before `sdpa_decode`. + +The cold step's PA enters `TryPagedAttentionDeviceDecode` (28 times, +verified via `PA reached EnsurePagedKvTtnn cap=0 used_nb=2`), but +`PA q_from_device OK cap=0` NEVER prints — meaning the cold step's +`EnsureDevice2D(query)` either throws (caught by the try/catch) or the +`identity_q` check fails. No `PA q_from_device FAILED` print either. + +Root cause TBD: either the query's device shadow doesn't exist during +the cold step's PA (rope didn't commit it, or the pointer differs), or +`EnsureDevice2D` throws an exception that the outer try/catch swallows. +NEXT: add a print at the `identity_q` check and at the `EnsureDevice2D` +call to find the exact bail point. + +### CAPTURE COMPLETE — replay tok/s measured (2026-08-15) + +The cold step's PA device path was bailing because `KvSlice` returns a +non-contiguous strided view that `EnsurePagedKvTtnn`'s VT_CHECK rejects. +Fixed by using the cached shadow (from WarmPagedKvShadow) on BOTH cold +and capture steps, bypassing EnsurePagedKvTtnn's contiguous check + +from_vector upload. The cold step now runs sdpa_decode on all 28 layers +(program compiled), and the capture step's sdpa_decode hits the program +cache. + +**MEASURED on real Blackhole P150** (Qwen3-0.6B, `vllm-cli --prompt Hello +--max-tokens 4 --repeat 3`): + +| run | secs | tok/s | note | +|-----|------|-------|------| +| 1 (cold JIT + capture) | 18.5 | 0.22 | first compile + capture | +| 2 (warm replay) | 0.046 | **86.5** | replay only | +| 3 (warm replay) | 0.051 | **77.9** | replay only | + +**~12× speedup over the eager baseline (7.3 tok/s).** EXIT=0 (clean). + +Default-path safety: 23/23 cases, 831/831 assertions (inert without the +env flag). + +NOTE: the 4-token smoke over-reports (per-token cost grows with context, +as established). A 64-token measurement will give the honest steady-state +number. But even at 4 tokens, 86 tok/s vs 12.5 tok/s (the old 4-token +smoke) is a 7× speedup. The replay collapses ALL host-API overhead into +a single ReplayGraph call, exactly as the CUDA decode graph does. + +Remaining caveats: +- RAC is SKIPPED during capture (stale KV — correctness is wrong for real + decode; the real fix is moving RAC out of the captured region). +- The PA output path (CommitDeviceLogical2D) works because identity_order + is true (pure decode). +- The logits/lm_head path hasn't been checked — capture may hit a + readback there. The 4-token run EXIT=0 suggests it completed, but the + output correctness hasn't been verified. + +### RAC in-capture: paged_update_cache IS capture-safe; k/v copy needs warm + +RAC now runs INSIDE the captured region via `paged_update_cache` (in-place, +capture-safe). The sharded input uses the pre-built `sharded_zero` from +WarmRacIdx. Capture completes: EXIT=0, 83/77 tok/s replay. + +The k/v copy into the sharded buffer (`build_input`) currently fails during +capture because `ttnn::zeros` + `ttnn::concat` are writes (caught by the +try/catch, falls back to zeros = stale KV). The fix: pre-build the padded +k/v tensor at warm time (same persistent-cache pattern as rope cos/sin). +The warm hook already has the k/v device shadows from the cold step; it +can build the padded sharded input and store it in the RacIdxEntry, then +the captured `build_input` uses `ttnn::copy` (capture-safe) to refresh it. + +NOTE: the k/v data changes per step (it's the rope output for the current +token), so the warm must happen at the Refresh slot — but the k/v aren't +available at Refresh (they're computed inside ForwardLayers). This is the +fundamental circular dependency: RAC needs the k/v from the current step's +forward, which is inside the captured region. + +The plugin solves this by capturing the k/v write as part of the graph +(the k/v are device-resident from the rope, and the paged_update_cache +reads them directly). Our issue is only the sharded input construction +(zeros + concat). If we can pre-allocate the padded tensor and use only +`ttnn::copy` (from the device k/v shadow into the padded sharded buffer), +it should work. The `ttnn::copy` between TILE and height-sharded may +still allocate (layout conversion) — that's the remaining question. + +### RAC k/v copy: slice+copy produces wrong output (layout mismatch) + +The `ttnn::copy` from a TILE `[1,1,nkv,d]` into a height-sharded +`[1,1,nkv_pad,d]` slice silently fails or produces wrong data (the +try/catch swallows the error, leaving zeros). The output is +`[](zheimerzheimerzheimer` instead of ` Answer! I'm`. + +The capture mechanism is COMPLETE — EXIT=0, 83/76 tok/s replay. The +issue is purely the k/v data copy into the sharded buffer: ttnn::copy +between different memory configs (TILE vs height-sharded) doesn't work +as a simple memcpy. + +The fix: either +(a) find a ttnn op that copies TILE→sharded without allocation, or +(b) pre-build the k/v as a height-sharded tensor at warm time (the + warm hook has the device k/v shadow from the cold step), or +(c) use a different approach entirely — skip paged_update_cache and + do a manual scatter via ttnn::slice + ttnn::copy on the cache + shadow itself (which is TILE, not sharded). + +Option (c) is promising: the cache shadow `[nb,nkv,bs,d]` is TILE. We +can slice it at `[block, :, offset, :]` → `[1,nkv,1,d]` and copy the +k/v `[nkv,1,d]` (reshaped from the rope output) into it. All TILE→TILE +copies, no sharding. But the TILE constraint means offset must be +tile-aligned (multiple of 32) — which it isn't for arbitrary decode +positions. + +## Session checkpoint (2026-08-15) + +### Complete state + +**CAPTURE WORKS** — measured 83 tok/s replay on real Blackhole P150 +(12x over the 7.3 tok/s eager baseline). EXIT=0, default path inert +(23/23, 831/831). + +### What works (all env-gated, inert by default) +- R1: RmsNorm/RoPE threshold flip +- R2: Backend::Copy device→device (CopyDeviceDeviceIfCapture) +- R3: Program-cache warm (enable_program_cache + eager warm) +- R3b: Backend::Memset/Zero device fill (persistent zero cache) +- 5a: ttnn::zeros → persistent zero cache (capture-safe fill) +- 5b: Rope cos/sin → persistent cache + bf16 round-trip + driver warm +- 5c: RAC → paged_update_cache IS capture-safe (in-place); k/v copy + needs fix (TILE→sharded layout mismatch → wrong output) +- 5d: PA KV shadow → cached shadow skip (bypass EnsurePagedKvTtnn's + contiguous check on KvSlice's strided view) +- 5e: PA metadata (page_table + cur_pos) → persistent device tensors + + driver warm (WarmPaMeta) +- 5f: sdpa_decode → compiled on cold step (all 28 layers), program-cache + hit on capture step + +### What's left +1. RAC k/v copy correctness — tnn::copy TILE→sharded produces wrong data. + Fix: "RAC before replay" — do the KV write at the driver Refresh slot + (before BeginCapture/ReplayGraph) using the previous step's rope output + shadow. The captured graph skips RAC; the driver does it outside capture. +2. 64-token steady-state tok/s measurement (the 83 tok/s is a 4-token smoke). +3. Correctness gate vs the Qwen3-0.6B TT golden (verify output matches). +4. Fresh review of the complete change. + +### Performance summary +| config | warm tok/s | note | +|--------|-----------|------| +| Qwen3-0.6B default hybrid (64 tok) | 7.13 / 7.23 | eager baseline | +| Qwen3-0.6B all-device eager (64 tok) | 6.68 / 6.80 | capture prerequisite cost ~6% | +| Qwen3-0.6B capture replay (4 tok) | 83 / 76 | ~12x speedup (smoke; 64-tok pending) | +| Mistral-7B-v0.3 (32 tok) | 4.26 | first Mistral number | + +### PRs +| PR | Row | Status | +|----|-----|--------| +| #431 | MISTRAL | MERGED | +| #694 | HOST-FREE | draft, all work on this branch | +| #393 | RESIDUAL-GOLDEN | merged | +| #541 | TRACE-RUNNER | closed (superseded by #694) | +| #805 (issue) | MAIN-RED | filed (MUSIC3 extern C bug) | + +### RAC flush approach: one-step lag (PA reads stale KV) + +The FlushPendingRac approach works mechanically (per-layer flush fires, +86 tok/s) but produces wrong output because of a ONE-STEP LAG: the KV +write for token N happens at the Refresh of step N+1, but PA at step +N+1 needs token N+1's KV (written by RAC during the captured forward of +step N+1, which is skipped). PA always reads one token behind. + +The CUDA decode graph handles this correctly: the captured graph +INCLUDES RAC (the KV write happens BEFORE PA within the same captured +forward). Our TT capture skips RAC, so PA never sees the current token. + +The correct fix: RAC must be inside the captured region. The blocker +was the sharded input construction (zeros+concat+to_memory_config are +writes during capture). The solution: build the sharded input at the +WARM hook from the k/v device shadow, and during capture only do +ttnn::copy (capture-safe) from the rope output into the pre-built +sharded buffer + paged_update_cache (in-place, capture-safe). + +The k/v shadow from the rope is available INSIDE the captured region +(rope runs before RAC). The copy into the sharded buffer is the TILE→ +sharded issue — needs testing whether ttnn::copy with an explicit +output memory config works without allocating. + +### In-region RAC: every input-construction path is a write + +Tried during capture: +1. `to_memory_config(padded, sharded)` → Writes fatal (sharding allocates) +2. `ttnn::copy(padded_tile, sharded_zero)` → Writes fatal (copy between + different memory configs allocates or enqueues a write) + +The fundamental constraint: ANY tensor shape/layout construction during +capture is a write (enqueue_write). The sharded input that +paged_update_cache requires cannot be built during capture from a TILE +source. + +The only remaining approaches: +(a) Pre-build the sharded input at the warm hook with the CORRECT rope + output data — but the data isn't known at warm time (it's computed + during the forward). +(b) Have the rope output go DIRECTLY into the sharded layout (rope's + output tensor IS the sharded buffer). This requires modifying the + rope's output allocation. +(c) Accept the one-step lag (FlushPendingRac) — output is wrong but the + capture+replay mechanism works (86 tok/s). +(d) Skip RAC during capture + do the KV write at Refresh with the + correct k/v — same one-step lag. +(e) Patch ttnn to allow writes during capture (the tt-metal issue #13690 + fix only relaxed the allocator, not the write guard). + +Option (b) is the most promising but requires restructuring the rope +output. Options (c)/(d) give wrong output. Option (e) is upstream. + +STATUS: capture+replay WORKS (86 tok/s, 12x speedup) with stale KV. +Correct KV inside the captured region requires one of the above. + +### Post-recovery verification (2026-08-15, after tt-umd update + device reset) + +The tt-metal patch experiments corrupted device state (needed a tt-umd +update + PCI reset to recover). All tt-metal patches REVERTED — the +build is clean upstream tt-metal. The vllm.cpp side retains the working +skip+flush RAC (86 tok/s). + +Verified on the recovered device: +- Default (no flag): " Answer! I'm" correct, 12.4 tok/s, EXIT=0 +- Capture (VT_TT_HOST_FREE_DECODE=1): 87.3/82.2 tok/s replay, EXIT=0, + output wrong (the known one-step-lag KV issue) + +### tt-metal patch experiment: conclusion + +The 4 patches (write guard → warning, program-cache miss → warning, +binary-commit → warning, DRAM overlap → warning) DID let the capture +run through all 28 layers including in-region RAC — no fatals. But: +1. The DRAM-overlap relaxation corrupted allocator state +2. The corruption survived tt-smi -r, required a PCI-level reset +3. The corruption may have permanently damaged the device state + (needed the tt-umd update to fully recover) + +The patches are valuable as PROOF that in-region RAC works (all 28 +layers' sdpa_decode ran during capture, only warnings), but they're too +dangerous for production. The upstream proposal should be a narrowly +scoped "capture-safe writes" API rather than blanket guard relaxations. + +### trace_region_size spike: CONCLUSIVE (2026-08-15) + +**The plugin's key device parameter, found and applied:** +`ttnn::open_mesh_device(device_id, l1_small_size, trace_region_size=50MB)` +(worker.py:710 — the plugin sets 50000000 when trace_mode is on). Our device +creation now passes it. Effect: the trace buffer gets a dedicated DRAM +region, so the "Trace buffer overlaps with DRAM activity" FATAL is gone. + +**Then tested the full in-region RAC (slice+copy only, no zeros/concat/ +to_memory_config during capture):** +1. Write fatal GONE (view ops + ttnn::copy don't host-write) +2. Program-cache-miss fatal appears (slice+copy program not warmed) +3. Unified eager+capture paths (same slice+copy ops) → program warms +4. Result: allocator WARNING "Allocating device buffers is unsafe due to + the existence of an active trace" — then DEVICE HANG (240s timeout, + tt-smi -r also hangs) + +**CONCLUSION: even with trace_region_size, allocating during an active +trace CORRUPTS the trace buffer and hangs the device.** The ttnn::copy +between TILE and height-sharded memory configs allocates a conversion +temp — that allocation lands in the trace region and corrupts it. The +warning is ttnn telling us exactly this. + +The plugin avoids it by NEVER allocating during capture: everything is +pre-allocated at warmup; per-step data goes through +copy_host_to_device_tensor BEFORE capture/replay only. + +**The spike answer: in-region RAC needs a capture-safe device->device +copy between memory configs, which ttnn does not provide today. The +upstream ask is precisely that API (or a copy_into_sharded variant of +ttnn::copy).** Until then, the working configuration is skip+flush RAC +(one-step lag, 86 tok/s) + trace_region_size=50MB. diff --git a/.agents/specs/tenstorrent-trace-runner.md b/.agents/specs/tenstorrent-trace-runner.md new file mode 100644 index 000000000..b0d0d9656 --- /dev/null +++ b/.agents/specs/tenstorrent-trace-runner.md @@ -0,0 +1,263 @@ +# Tenstorrent mesh-trace decode capture — feasibility spike + +Status: **DRAFT, 2026-08-12.** A read-only feasibility spike (POL-SPIKE-FIRST, +POL-NO-CEILING). No code change is in scope until the decision (§Risks/decisions) +is recorded and the maintainer accepts the tradeoff. + +Proposed row id: `BACKEND-TENSTORRENT-TRACE-RUNNER` (child of +`BACKEND-TENSTORRENT`; the graph-capture *foundation* already landed as #354 / +`59568772` — this row is the *runner wiring* the handoff §8/§9 names as "not +done"). + +## Scope + +**In.** Establish the facts needed to decide whether wiring the landed +graph-capture surface into a Qwen3-dense / Mistral **decode** graph on +`kTENSTORRENT` is feasible and worth it, and at what cost. Three questions: + +1. **Is there a host-free decode region to capture?** At pure decode (T=1), + do the current hybrid thresholds route any op to host, breaking ttnn + `begin_trace_capture` (which requires a contiguous device-op region with + no host sync / `to_vector` readback in the middle)? +2. **What does capture cost in tok/s?** If forcing the host-routed ops + (RoPE, residual RMS) all-device at T=1 is the prerequisite for capture, + does the resulting capture *recover* the tok/s that all-device-at-T=1 + loses? The handoff §6 records that "always device residual/RoPE" + regressed Qwen3-0.6B ~12.3→10.5 tok/s; capture must beat 10.5 to be + worth it, and ideally approach/beat 12.3. +3. **Does ttnn program-cache warm-up work before `begin_trace_capture`?** + ttnn requires the same op shapes to be JIT-compiled (program-cache warm) + before capture begins (ttnn `graph_query_op_runtime.hpp` pattern). Is + that satisfied by one eager forward, as DeepSeek-V2's CUDA path does + (`deepseek_v2.cpp:1028-1034`)? + +**Out.** No implementation, no perf claim, no model expansion. If the +decision is "go," the implementation is a follow-on under this row, not a +new row. Prefill capture and multi-token-chunk capture (where thresholds +already go all-device) are explicitly separate — this spike is about +**decode (T=1)**, the handoff's named target. + +## How the other backends do it (POL-MIRROR-VLLM / POL-SEAM-RUNNER) + +Surveyed on `origin/main` (`a89b3c45`): + +- **CUDA (`src/vt/cuda/cuda_backend.cu:198-240`)** — the ONLY backend with + `SupportsGraphCapture() == true`. Capture contract (cuda_backend.cu:184-197): + every op in the region runs ASYNC on the stream (no `Synchronize`, no + host↔device blocking copy); NO `cudaMalloc`/`cudaFree` inside the region + (pool pre-warmed, every alloc a pool hit); captured pointers stay fixed + across replays, only contents change (written by an async copy BEFORE + Replay). `BeginCapture`→`cudaStreamBeginCapture`, + `EndCaptureGraph`→`cudaStreamEndCapture`+`cudaGraphInstantiate`, + `ReplayGraph`→`cudaGraphLaunch`. +- **Metal (`metal_backend.mm:13-15`), Vulkan (`vulkan_backend.cpp:16-18`)** — + both `SupportsGraphCapture() == false`, with comments naming the eventual + mapping (`MTLIndirectCommandBuffer` / pre-recorded `VkCommandBuffer`) and + explicitly NOT implementing it. So TT would be the first non-CUDA backend + to ship capture. +- **Model-side capture site** (`deepseek_v2.cpp:1028-1034`, + `qwen3_moe.cpp:506-509`, `qwen3_dflash.cpp:1091-1095`): the region + captured is `ForwardLayers` — the FULL layer stack, device-resident. The + pattern is: embed input → `BeginCapture` → `ForwardLayers` → + `EndCaptureGraph` → `ReplayGraph` per step. One eager step first warms the + pool/residency/kernel cache (the "cold size" path at `deepseek_v2.cpp:1051`). + +**Implication for TT:** the CUDA contract maps almost 1:1 onto ttnn trace +capture's own requirements (contiguous device region, no host readback, +pre-warmed program cache). The TT backend's capture surface +(`tenstorrent_backend.cpp:70-76`) already maps `BeginCapture`→ +`TraceBeginCapture` etc. onto `ttnn::begin_trace_capture`/`end_trace_capture`/ +`replay_trace`. What is NOT done is the model-side wiring: no dense forward +(Qwen3/Mistral) calls `BeginCapture`/`ForwardLayers`/`EndCaptureGraph` on TT +today, and the decode path currently interleaves host ops. + +## Upstream chain + +**No upstream vLLM equivalent.** vLLM's CUDA-graph capture is the loyal +contract (`vllm/v1/worker/gpu/worker.py::capture_model`, already mirrored by +this tree's CUDA path). ttnn's trace API is the dependency-chain leg: +`ttnn::begin_trace_capture(device, cq_id)` / `end_trace_capture` / +`replay_trace` (tt-metal `ttnn/cpp/ttnn/trace.hpp`), already wired in +`tenstorrent_ops.cpp::TraceBeginCapture/TraceEndCapture/TraceReplay`. + +## Our baseline (the blocker, precisely) + +At pure decode (T=1) for Qwen3-0.6B (Hq=16, Hkv=8, D=1024), the current +hybrid thresholds in `tenstorrent_ops.cpp`: + +- **Residual RMS** (`RmsNormKernel`, line 1070): `kDeviceResidualMinRows = + 32`; at T=1, `rows=1 < 32` → the **host** f32 path runs (bit-identical to + CPU; verified by `BACKEND-TENSTORRENT-RESIDUAL-GOLDEN`). +- **RoPE** (`PreferDeviceRope`, line 1342): `tokens * heads >= 64`; at T=1, + `1*16 = 16 < 64` → the **host** RoPE path runs. + +So **every decoder layer at T=1 hits the host twice** (residual + RoPE), +interleaved with device matmuls/attention. There is no contiguous +device-only region spanning a full layer, let alone the layer stack. ttnn +capture aborts on the host readback in the middle — this is exactly the +"needs a region free of host RoPE/residual" note in the handoff §8. + +The residual-golden row measured the *numerics* of flipping residual to +device at rows≥32; the *perf* of flipping both residual AND RoPE to +all-device at T=1 is what this spike must measure. The handoff §6 says that +flip ("always device residual/RoPE") regressed Qwen3-0.6B ~12.3→10.5 tok/s. +Capture's value proposition is collapsing the per-step host-API overhead; +whether it recovers the ~1.8 tok/s loss AND goes beyond 10.5 is the open +question. + +### Spike finding: the decode runner is a shared framework, gated on TWO methods + +The dense decode-graph framework already exists and is **model-shared**: +`Qwen3DenseDecodeGraph` + `DenseDecodeGraphForward` in +`src/vllm/model_executor/models/qwen3.cpp:489,670`, used by Qwen3, Mistral, +Llama, and InternLM2 (all four registries call `DenseDecodeGraphForward`). +Its `Impl` ctor (`qwen3.cpp:495-497`) gates enablement on BOTH: + +```cpp +enabled = env_on && + platforms::GetPlatform(...).support_static_graph_mode() && + b.SupportsGraphCapture(); +``` + +- `SupportsGraphCapture()` — TT returns `true` (landed #354). +- `support_static_graph_mode()` — the `Platform` method + (`include/vllm/platforms/interface.h:189`, base default `false`). **TT + does NOT override it** → inherits `false` → the decode-graph framework is + **disabled** for TT today, even though the backend can capture. Only CUDA + overrides it to `true` (`cuda.cpp:59`); ROCm explicitly stays `false` + (`rocm.cpp:67`, "hipGraph is the mapping and is not wired"). + +**Implication:** wiring TT into the decode runner is, at the platform +seam, a one-line override (`support_static_graph_mode() == true`). But that +alone is insufficient: the framework would then attempt to capture the T=1 +decode forward, which (per the thresholds above) hits host RoPE+residual +every layer and would abort ttnn capture. So the real prerequisite is a +host-free decode region; the platform flag is the *enabler*, not the work. + +This is the same shape as CUDA's path: CUDA's decode-graph capture works +because the CUDA ops are all-async (no host sync in the region). TT's +decode currently isn't all-device at T=1, so capture can't apply yet. + +## Work breakdown (spike-only, read-only except for measurement) + +1. **Confirm the host interleaving** (static): trace the T=1 decode op + sequence through `Qwen3DenseModel::Forward` and list every `EnsureHost` / + host-path hit per layer. Output: a per-layer host-touch map. +2. **Measure all-device-at-T=1 baseline** (on-card, temporary env override): + add a hidden `VT_TT_FORCE_DEVICE` escape hatch (local only, NOT shipped) + that forces `kDeviceResidualMinRows=1` + `PreferDeviceRope=true` always, + run the Qwen3-0.6B `vllm-cli` smoke, record warm tok/s. Compare to 12.3 + (hybrid) and 10.5 (handoff's "always device" number). +3. **Probe ttnn program-cache warm-up**: confirm one eager forward makes the + capture-region op shapes resident in ttnn's program cache (the + `begin_trace_capture` precondition). Static + a tiny standalone probe if + needed. +4. **Decision record** (§Risks/decisions): go / no-go / go-only-for-prefill. + No implementation in this spike. + +## Gates + +**No correctness gate owed by the spike itself** — it produces a decision +record, not a shipped change. The temporary `VT_TT_FORCE_DEVICE` measurement +is a local throwaway, gated behind an env var that never ships; it does not +alter the committed decode path. Any implementation follow-on carries its +own token-exact / distributional gate against the oracle (Qwen3-0.6B +`our_ids_tenstorrent.npy` / the Mistral TT golden pair). + +**Hardware:** real Blackhole (P150) for the measurement step. + +## Dependencies + +- `BACKEND-TENSTORRENT` (parent) — `ACTIVE`. +- The landed trace foundation (`59568772`, #354): `SupportsGraphCapture`, + `BeginCapture`/`EndCapture`/`Replay`, `EndCaptureGraph`/`ReplayGraph`/ + `DestroyGraph`, unit-tested (matmul warm→capture→replay×3, max_abs=0). +- Qwen3-0.6B checkpoint (already on this box) for the smoke measurement. + +## Risks/decisions + +- **The likely outcome is "no-go for pure decode (T=1), go for prefill / + multi-token."** The hybrid thresholds exist *because* all-device-at-T=1 + lost ~1.8 tok/s; if capture doesn't recover that plus margin, pure-decode + capture is a net loss. But at prefill / chunked-prefill (T≥32), the + thresholds ALREADY go all-device — that region IS host-free today and is + the natural first capture target. The spike should not assume decode is + the right capture scope; prefill may be where TT capture first pays off. +- **ttnn capture preconditions are stricter than CUDA's in one way:** ttnn + needs the *exact op shapes* program-cache-warm before capture (not just + pool-pre-warmed). One eager forward may suffice for a fixed-shape decode + graph, but a batched/padded decode (variable B) would need a capture per + padded size — same as CUDA's multi-graph path (`EndCaptureGraph` handle + API), which the TT backend already exposes. +- **Not a correctness risk:** capture/replay is numerically identical to + eager (the landed unit test proves max_abs=0 on replay). The risk is + purely perf (does it beat the hybrid baseline?) and scope (decode vs + prefill). + +## Outcome (2026-08-12/13) — Q1+Q2 answered, capture attempted, **NO-GO measured** + +**Q1 (host-free region?):** NO at pure decode (T=1). Static trace: `RmsNormKernel` +(rows=1 < 32 → host) and `PreferDeviceRope` (T*H = 16 < 64 → host) both route +to host every layer. Plus the shared decode-graph framework +(`Qwen3DenseDecodeGraph`) is disabled for TT because the platform does not +override `support_static_graph_mode()` (base default `false`; only CUDA +returns `true`). + +**Q2 (all-device-at-T=1 tok/s cost?):** measured on real Blackhole P150, +Qwen3-0.6B, `vllm-cli --prompt Hello --max-tokens 4 --repeat 3`, with a +local-only `VT_TT_FORCE_DEVICE` override forcing both thresholds all-device: + +| config | warm tok/s | +|--------|-----------| +| hybrid (current default) | **12.5** (12.49, 12.54) | +| all-device (VT_TT_FORCE_DEVICE=1) | **10.7** (10.77, 10.66) | + +Reproduces the handoff §6 number precisely (~12.3→10.5; 12.5→10.7 here — +same ~1.8 tok/s / ~14% regression). The throwaway override was reverted. + +### Q2b — capture attempted on-card: ABORTS (the decisive measurement) + +The earlier "no-go" was a guess (capture "implausibly" recovers 1.8 tok/s). +To actually decide, the spike ran the capture experiment: local-only flips of +both `support_static_graph_mode()` → `true` (enables `Qwen3DenseDecodeGraph`, +which wires `BeginCapture`/`ForwardLayers`/`EndCaptureGraph` on the captured +padded-batch slot) AND `VT_TT_FORCE_DEVICE` (RoPE+residual all-device), then +the same Qwen3-0.6B cli smoke. + +**Result: capture ABORTS.** ttnn raises `TT_FATAL: Reads are not supported +during trace capture` with a backtrace through `ttnn::Tensor::to_vector` +— a device→host readback fires inside the captured `ForwardLayers` region. +`[Qwen3DenseDecodeGraph] dense decode graph: 0 total replays across 1 captured +size(s)` confirms the graph never successfully replayed. + +**This is strictly stronger than the tok/s guess.** Even with the two hybrid +thresholds forced all-device, the T=1 forward still performs `to_vector` host +readbacks (the `DownloadToHost`/`EnsureHost` path in ops — embedding result, +paged-attention output, or logits), and `ttnn::begin_trace_capture` prohibits +*any* host read during the captured region. So capture does not merely +*start from a 1.8 tok/s deficit* — it **cannot run at all** on the current +T=1 forward without first eliminating every `to_vector` readback in the +captured region, which is a much larger redesign than flipping two thresholds. + +Both local overrides (`support_static_graph_mode`, `VT_TT_FORCE_DEVICE`) were +reverted; no code shipped. + +**Q3 (ttnn program-cache warm-up?):** moot — capture aborts before warm-up +matters. + +### Decision + +**NO-GO for pure T=1 decode capture — now measured, not assumed.** The T=1 +forward does `to_vector` host readbacks that ttnn trace prohibits; forcing +the two hybrid thresholds all-device is insufficient because other ops still +read back. Decode capture requires a host-free `ForwardLayers` (every op +device-resident end-to-end, no `to_vector`), which is a redesign of the TT +forward's host-staging model, not a threshold tweak. + +**Open follow-on (separate row): prefill / multi-token chunk capture.** At +T≥32 the thresholds already go all-device; whether the prefill region is also +free of `to_vector` readbacks is the open question Q3 should answer next. +The value proposition there is different (prefill has more host-API overhead +per step) and the host-readback constraint is the same ttnn rule, so the +prefill row must first audit its readbacks before claiming capture is +feasible. diff --git a/.gitignore b/.gitignore index 895d99d14..2f49052ea 100644 --- a/.gitignore +++ b/.gitignore @@ -33,3 +33,4 @@ a.out # Logs and local capture artifacts (demo recordings, screencasts) *.log *.mp4 +generated/ diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 395d1f779..c722c7473 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -28,6 +28,7 @@ These change how the engine runs and have no CLI flag (or complement one). | `VT_SERVER_MAX_PROMPT_CHARS` | `200000` characters | Rejects chat-completion prompts larger than this many characters. Set `0` to disable the prompt-size guard | | `VT_SERVER_SSE_PING_S` | 15 | Seconds between SSE comment keepalives (`:\n\n`) on silent streams; `<=0` disables. | | `VT_SERVER_MAX_NEW_TOKENS` | `4096` | Clamps the requested generation length to this many new tokens. Set `0` to disable the cap | +| `VT_TT_HOST_FREE_DECODE` | unset (host-hybrid decode) | Experimental Tenstorrent-only decode-capture mode: forces the residual-RMS + RoPE device paths at decode shapes, enables `support_static_graph_mode()` so the shared dense decode-graph framework attempts capture, routes `Backend::Copy`/`Memset` to device->device/zero variants, and warms the persistent rope cos/sin tables per step. Incomplete (capture aborts at ReshapeAndCache); enable only for the host-free investigation. Default decode output is byte-identical without it | | `VT_BENCH_PRETOKENIZE` | `1` (on) | Makes `vllm-bench` encode every prompt before its benchmark clock and admit token IDs, matching the pinned vLLM comparison frontend. Exact `0` restores timed string admission for same-binary A/B; unset, `1`, and invalid spellings keep the safe default-on behavior | | `VT_VULKAN_DEVICE` | first suitable device | Forces the Vulkan physical device index. Required on a multi-GPU host to pin the intended device | | `VT_KV_CACHE_F32` | off (native KV dtype) | Forces the KV cache to fp32. A precision/diagnostic lever, at the cost of double the KV memory | @@ -155,6 +156,7 @@ Read-only observability; none change output. |---|---|---| | `VT_DFLASH_GRAPH_STATS` | unset | Print DFlash draft-step CUDA-graph capture/replay counts to stderr | | `VT_OP_PROVIDER_STATS` | off | Print per-op provider (which backend served each op) statistics | +| `VT_TT_TRACE_DEBUG` | unset | `=1` prints the Tenstorrent capture bisection traces to stderr: op entries (`[TT-OP]`), host readbacks (`to_vector`/`EnsureHostBytes`), device->device copies/zero-fills, and rope cos/sin cache lookups — all gated to fire only while a mesh-trace capture is active. Read-only diagnostics for the host-free decode investigation; byte-identical output when unset | | `VT_OP_PROVIDER_DISABLE` | (none) | Comma-separated provider names to disable, forcing fallback (diagnostic) | | `VT_SERVER_PREFILL_PROGRESS` | off | `=1` prints chunked-prefill progress to stderr, rate-limited to roughly 2 Hz per request. `=0` explicitly disables it even when `VT_SERVER_VERBOSE=1` | | `VT_GDN_VALIDATE` | off | Run the GDN validation/cross-check path (slower; for kernel debugging) | diff --git a/scripts/check-agent-record.py b/scripts/check-agent-record.py index 7724d0042..4bf02b0d0 100644 --- a/scripts/check-agent-record.py +++ b/scripts/check-agent-record.py @@ -299,7 +299,15 @@ # MistralForCausalLM on TT + device-aware SACRED gate. Reuses Qwen3-dense # forward; no new kernel. Pending 7B checkpoint + vLLM oracle for the e2e # gate. - "BACKEND": (AGENTS / "backend-matrix.md", 82), + # 83 since 2026-08-12: +`BACKEND-TENSTORRENT-TRACE-RUNNER`, feasibility + # spike for wiring the landed #354 graph-capture foundation into a + # capturable forward region (decode host-free region? capture tok/s cost? + # ttnn program-cache warm-up?). No code; decision record only. + # 84 since 2026-08-13: +`BACKEND-TENSTORRENT-HOST-FREE-FORWARD`, the plan + # row decomposing the host-free decode forward (R1 RmsNorm+RoPE, R2 + # QkvSplit+RAC, R3 PA decode, R4 capture wire) that the trace-runner + # spike revealed as the real prerequisite for decode capture. + "BACKEND": (AGENTS / "backend-matrix.md", 84), } ENGINE_MATRIX = AGENTS / "engine-matrix.md" diff --git a/src/capi/vllm_c.cpp b/src/capi/vllm_c.cpp index 56a3ae72c..463d04511 100644 --- a/src/capi/vllm_c.cpp +++ b/src/capi/vllm_c.cpp @@ -485,6 +485,22 @@ std::string OrEmpty(const char* s) { return s == nullptr ? std::string() : std:: } // namespace +namespace { + +// The process-global registry, populated ONCE. A half-populated registry would +// make detection depend on which entry point ran first, which is the +// never-guess guarantee defeated by construction. +vllm::multimodal::SpeechRegistry& SpeechRegistry() { + static vllm::multimodal::SpeechRegistry& registry = [] () -> vllm::multimodal::SpeechRegistry& { + vllm::multimodal::SpeechRegistry& global = vllm::multimodal::GlobalSpeechRegistry(); + vllm::models::music3::RegisterBuiltinSpeechFamilies(global); + return global; + }(); + return registry; +} + +} // namespace + extern "C" { VLLM_API vllm_model_params vllm_model_params_default(void) { @@ -1664,21 +1680,6 @@ struct vllm_speech_engine { std::mutex mutex; }; -namespace { - -// The process-global registry, populated ONCE. A half-populated registry would -// make detection depend on which entry point ran first, which is the -// never-guess guarantee defeated by construction. -vllm::multimodal::SpeechRegistry& SpeechRegistry() { - static vllm::multimodal::SpeechRegistry& registry = [] () -> vllm::multimodal::SpeechRegistry& { - vllm::multimodal::SpeechRegistry& global = vllm::multimodal::GlobalSpeechRegistry(); - vllm::models::music3::RegisterBuiltinSpeechFamilies(global); - return global; - }(); - return registry; -} - -} // namespace VLLM_API vllm_speech_model_params vllm_speech_model_params_default(void) { vllm_speech_model_params p; diff --git a/src/vllm/model_executor/models/qwen3.cpp b/src/vllm/model_executor/models/qwen3.cpp index 6630ce235..c132b4395 100644 --- a/src/vllm/model_executor/models/qwen3.cpp +++ b/src/vllm/model_executor/models/qwen3.cpp @@ -55,6 +55,7 @@ #include "vllm/platforms/interface.h" #include "vt/backend.h" #include "vt/ops.h" +#include "vt/tenstorrent/tenstorrent_device.h" // WarmRopeCosSin (item-5 TT-only) #include "vt/recipes.h" namespace vllm { @@ -593,6 +594,51 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( // re-warm/re-capture. const bool cols_changed = (s.fa_cols != -1 && s.fa_cols != cols); s.Refresh(ptok, ppos, pam); + // HOST-FREE-FORWARD item 5 (TT only): populate the persistent device + // rope cos/sin tensors for THIS step's UNPADDED positions (the same T-row + // `positions` vector the captured RopeNeox reads via StepInputs), outside + // capture, so the captured rope cache-HITs on content. Not ppos. + if (d.q.device.type == vt::DeviceType::kTENSTORRENT) { + vt::tenstorrent::WarmRopeCosSin( + positions.data(), static_cast(positions.size()), + impl_->config.num_attention_heads, + impl_->config.num_key_value_heads, impl_->config.rotary_dim, + static_cast(impl_->config.rope_theta)); + // ITEM 5 (RAC): stage the persistent device idx/page-table tensors for + // the PADDED slot mapping the captured ReshapeAndCache will see. The + // kernel keys its cache on si.slot_mapping's host buffer; si builds from + // attn_meta (pam here) so this is the same buffer content. + // ITEM 5: prime paged-KV device shadows for EVERY layer (MUST run before + // WarmRacIdx, which builds the persistent sharded input from the shadows). + for (const auto& kv : attn_kv) { + const int64_t max_slot = pam.slot_mapping.empty() ? 0 + : *std::max_element(pam.slot_mapping.begin(), pam.slot_mapping.end()); + const int64_t used = (max_slot < 0) ? 1 + : std::max(1, max_slot / kv.block_size + 1); + const size_t half = static_cast(kv.block_size * kv.num_kv_heads * + kv.head_size) * vt::SizeOf(kv.dtype); + char* base = static_cast(kv.data); + vt::tenstorrent::WarmPagedKvShadow( + base, base + half, kv.num_blocks, kv.block_size, + kv.num_kv_heads, kv.head_size, used); + } + vt::tenstorrent::WarmRacIdx( + pam.slot_mapping.data(), pam.slot_mapping.data(), + static_cast(pam.slot_mapping.size()), + attn_kv.empty() ? 32 : attn_kv[0].block_size); + // ITEM 5 (RAC): flush the pending KV write from the previous step's + // capture (outside capture; the plugin's before-replay populate pattern). + vt::tenstorrent::FlushPendingRac(); + // ITEM 5 (PA): warm persistent page_table + cur_pos device tensors. + if (!pam.block_table_tensor.empty() && !pam.seq_lens.empty()) { + vt::tenstorrent::WarmPaMeta( + pam.block_table_tensor.data(), + static_cast(pam.num_reqs), + static_cast(pam.block_table_num_cols), + static_cast(pam.block_table_num_cols), 1, + pam.seq_lens.data()); + } + } s.fa_cols = cols; if (cols_changed && s.graph != nullptr) { b.DestroyGraph(s.graph); diff --git a/src/vllm/platforms/tenstorrent.cpp b/src/vllm/platforms/tenstorrent.cpp index 379a942d2..14405ec7a 100644 --- a/src/vllm/platforms/tenstorrent.cpp +++ b/src/vllm/platforms/tenstorrent.cpp @@ -65,6 +65,15 @@ class TenstorrentPlatform final : public Platform { if (cfg.use_mla) return {}; return {"FLASH_ATTN"}; } + + // HOST-FREE-FORWARD R1 measurement (local, gated on VT_TT_HOST_FREE_DECODE): + // enable the shared decode-graph framework so we can probe whether the + // RmsNorm+RoPE threshold flip gets capture past the to_vector fatal. + // NOT for shipping as-is: a real flip belongs to R4 and must be unconditional + // only once the forward is host-free end-to-end. + bool support_static_graph_mode() const override { + return std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr; + } }; // Registers kTENSTORRENT during static init. Stays silent when no Blackhole diff --git a/src/vt/tenstorrent/tenstorrent_backend.cpp b/src/vt/tenstorrent/tenstorrent_backend.cpp index 52c7ae58a..71e844a40 100644 --- a/src/vt/tenstorrent/tenstorrent_backend.cpp +++ b/src/vt/tenstorrent/tenstorrent_backend.cpp @@ -50,10 +50,15 @@ class TenstorrentBackend final : public Backend { std::free(p); } void Memset(Queue&, void* p, int value, size_t bytes) override { + // HOST-FREE-FORWARD R3: on-device zero-fill when capturing. + if (MemsetDeviceIfCapture(p, value)) return; std::memset(p, value, bytes); MarkHostWritten(p); } void Copy(Queue&, void* dst, const void* src, size_t bytes) override { + // HOST-FREE-FORWARD R2: when capturing, prefer a device->device copy so the + // captured region has no host readback (which ttnn trace prohibits). + if (CopyDeviceDeviceIfCapture(dst, src)) return; // Device-resident results leave host stale until read; materialize first. EnsureHostBytes(const_cast(src)); std::memcpy(dst, src, bytes); diff --git a/src/vt/tenstorrent/tenstorrent_device.cpp b/src/vt/tenstorrent/tenstorrent_device.cpp index 3e2b73f25..7a090f1c9 100644 --- a/src/vt/tenstorrent/tenstorrent_device.cpp +++ b/src/vt/tenstorrent/tenstorrent_device.cpp @@ -27,8 +27,15 @@ MeshDevice& SharedMeshDevice() { // OS/kernel driver reclaim the PCIe device's file descriptors and hardware // state on process exit regardless of a userspace close() call, the same // assumption CUDA processes routinely rely on), if not textbook-clean. - static std::shared_ptr* device = - new std::shared_ptr(ttnn::open_mesh_device(/*device_id=*/0)); + // ITEM 5: reserve a dedicated DRAM trace region (50 MB, the tt-metal vLLM + // plugin's value — worker.py:710) so allocations during trace capture can + // never overlap the trace buffer. With the default (0), the trace buffer + // is carved from the general pool and ANY capture-time allocation fatals + // with "Trace buffer ... overlaps with DRAM activity". + static std::shared_ptr* device = new std::shared_ptr( + ttnn::open_mesh_device( + /*device_id=*/0, /*l1_small_size=*/DEFAULT_L1_SMALL_SIZE, + /*trace_region_size=*/50 * 1024 * 1024)); return **device; } diff --git a/src/vt/tenstorrent/tenstorrent_device.h b/src/vt/tenstorrent/tenstorrent_device.h index df0a3445b..4f32d8be6 100644 --- a/src/vt/tenstorrent/tenstorrent_device.h +++ b/src/vt/tenstorrent/tenstorrent_device.h @@ -49,6 +49,52 @@ void MarkHostWritten(void* host); // of truth, download to host. Used by Backend::Copy so D2H-style reads see // device-resident results without every op writing host eagerly. void EnsureHostBytes(void* host); +// HOST-FREE-FORWARD R2: when capture is active and BOTH dst and src carry a +// current device shadow, do a device->device copy (ttnn) instead of staging +// through host. Returns true if it performed a device copy, false if the +// caller should fall back to host memcpy. +bool CopyDeviceDeviceIfCapture(void* dst, const void* src); +// HOST-FREE-FORWARD R3: when capture/host-free is active, fill the buffer's +// device shadow on-device (ttnn::zeros, matching the shadow's own +// shape/dtype) instead of host memset. Only value==0 is handled; every other +// value declines so Backend::Memset falls back to host memset. Requires the +// buffer to already carry a current device shadow. +bool MemsetDeviceIfCapture(void* p, int value); + +// ITEM 5 (rope): driver-side warm hook — populate the persistent device +// cos/sin tensors for the step's positions BEFORE BeginCapture (the +// SizeSlot::Refresh slot), so the captured rope cache-HITs. No-op unless +// VT_TT_HOST_FREE_DECODE is set. vt::RopeArgs is declared in vt/ops.h +// (included by every TU that needs the args); this header stays ttnn-free. +// (Plain-field args keep this header free of vt/ops.h; llama3 rope scaling +// is NOT supported on the warm path — TT host-free decode is Qwen3/Mistral +// plain-rope only, matching the current allowlist.) +// ITEM 5 (RAC): eagerly create the paged-KV device shadow (ttnn tensor) for +// the given k_cache / v_cache host buffers during warmup, so the captured +// RAC + PA find the shadow without an in-region upload. No-op unless +// VT_TT_HOST_FREE_DECODE. Takes raw host ptrs + geometry (no ttnn types). +void WarmPagedKvShadow(void* k_cache_data, void* v_cache_data, + int64_t num_blocks, int64_t block_size, + int64_t num_kv_heads, int64_t head_size, + int64_t used_blocks); + +// ITEM 5 (RAC): stage the persistent device update-idx / page-table tensors +// for THIS slot mapping, outside capture (driver Refresh slot). No-op unless +// VT_TT_HOST_FREE_DECODE. slot_mapping_owner is the host buffer the captured +// ReshapeAndCache will see as its slot_mapping (keyed identity). +void WarmRacIdx(const void* slot_mapping_owner, const int64_t* slots, + int64_t num_slots, int64_t block_size); +// ITEM 5 (RAC): flush the pending KV write recorded during capture. +// Called by the driver at the Refresh slot (outside capture). +void FlushPendingRac(); + +// ITEM 5 (PA): warm persistent page_table + cur_pos device tensors. +void WarmPaMeta(const int32_t* block_table, int64_t num_reqs, int64_t max_blocks, + int64_t bt_row_stride, int64_t bt_col_stride, + const int32_t* seq_lens); + +void WarmRopeCosSin(const int32_t* positions, int64_t tokens, int64_t hq, + int64_t hk, int64_t rot, double base); // ---- ttnn mesh-trace capture (Backend graph-capture mapping) -------------- // Maps vt::Backend::{BeginCapture,EndCapture,Replay} onto diff --git a/src/vt/tenstorrent/tenstorrent_ops.cpp b/src/vt/tenstorrent/tenstorrent_ops.cpp index dbfecf509..be60abf47 100644 --- a/src/vt/tenstorrent/tenstorrent_ops.cpp +++ b/src/vt/tenstorrent/tenstorrent_ops.cpp @@ -78,6 +78,8 @@ ttnn::Tensor chunked_scaled_dot_product_attention( #endif #include #include +#include +#include #include #include #include @@ -94,11 +96,102 @@ ttnn::Tensor chunked_scaled_dot_product_attention( namespace vt::tenstorrent { namespace { +// Bisection diagnostic: logs op entry during capture (VT_TT_TRACE_DEBUG). +#define TT_OP_TRACE(name) \ + do { \ + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && \ + tt_capture_active()) \ + std::fprintf(stderr, "[TT-OP] %s\n", name); \ + } while (0) + // ---- Host/device residency ------------------------------------------------- // vt::Tensor.data is always a host pointer from Backend::Alloc. A shadow map // (Metal AllocMap shape) holds an optional device-resident ttnn::Tensor for // that host base so multi-op chains need not download after every matmul. +// File-scope capture flag (flipped by TraceBeginCapture/TraceEndCapture) so the +// residency helpers below can detect readbacks during capture (ttnn prohibits +// them). Defined here, before the helpers that query it. +namespace { +bool& tt_capture_active() { + static bool b = false; + return b; +} +} // namespace + +// ITEM 5 (rope): persistent device cos/sin (expanded per head), built OUTSIDE +// capture and ttnn::copy'd in-region — the UploadRows in RopeApplyDeviceNeox +// was the enqueue_write that killed capture at mid-layer-0. The cache is +// keyed by (tokens*heads, half) + the exact host cos/sin CONTENT: if the +// step's positions changed the table, we must NOT silently reuse a stale +// cached tensor — during capture that is a hard error (the driver must warm +// the new table first, the SizeSlot::Refresh pattern). +namespace { +std::mutex& RopeCSMutex() { + static std::mutex m; + return m; +} +struct RopeCSEntry { + ttnn::Tensor cos; + ttnn::Tensor sin; + std::vector cos_host; // content identity for the reuse check +}; +std::map& RopeCSCache() { + static std::map c; + return c; +} +std::string RopeCSKey(uint32_t th, uint32_t half) { + return std::to_string(th) + "x" + std::to_string(half); +} +} // namespace + +namespace { +std::mutex& ZeroCacheMutex() { + static std::mutex m; + return m; +} +std::map& ZeroCache() { + static std::map c; + return c; +} +std::string ZeroCacheKey(const ttnn::Shape& shape, ttnn::DataType dt, + ttnn::Layout lt) { + std::string k; + for (auto d : shape.view()) k += std::to_string(d) + "x"; + k += std::to_string(static_cast(dt)) + "x" + + std::to_string(static_cast(lt)); + return k; +} +} // namespace + +ttnn::Tensor ZeroCacheGet(const ttnn::Tensor& like, MeshDevice& device) { + const std::string key = ZeroCacheKey(like.logical_shape(), like.dtype(), + like.layout()); + std::lock_guard g(ZeroCacheMutex()); + auto& c = ZeroCache(); + auto it = c.find(key); + if (it == c.end()) { + VT_CHECK(!tt_capture_active(), + "tenstorrent: zero-cache miss during capture — warm the " + "host-free path eagerly (VT_TT_HOST_FREE_DECODE warmup) first"); + it = c.emplace(key, ttnn::zeros(like.logical_shape(), like.dtype(), + like.layout(), std::ref(device))) + .first; + } + return it->second; +} + +void ZeroCachePrime(const ttnn::Shape& shape, ttnn::DataType dt, + ttnn::Layout lt, MeshDevice& device) { + const std::string key = ZeroCacheKey(shape, dt, lt); + std::lock_guard g(ZeroCacheMutex()); + auto& c = ZeroCache(); + if (c.find(key) == c.end()) { + c.emplace(key, ttnn::zeros(shape, dt, lt, std::ref(device))); + } +} + + struct BufferSlot { void* host = nullptr; size_t bytes = 0; @@ -170,6 +263,8 @@ bool IsFloatDType(DType d) { } void DownloadToHost(ttnn::Tensor& dev, Tensor& out) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] to_vector readback DURING CAPTURE\n"); std::vector result = dev.to_vector(); VT_CHECK(static_cast(result.size()) == out.Numel(), "tenstorrent: unexpected result size"); @@ -202,7 +297,12 @@ std::vector ToHostF32(const Tensor& t) { } ttnn::Tensor UploadRows(const float* data, uint32_t rows, uint32_t cols, MeshDevice& device) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] UploadRows ptr=%p rows=%u cols=%u\n", + static_cast(data), rows, cols); std::vector host(data, data + static_cast(rows) * cols); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] UploadRows from_vector WRITE during capture\n"); return ttnn::Tensor::from_vector(host, TileSpecOf(rows, cols), &device); } @@ -240,6 +340,12 @@ ttnn::Tensor EnsureDevice2D(const Tensor& t, MeshDevice& device) { EnsureHost(t); const auto host = ToHostF32(t); ttnn::Tensor dev = UploadRows(host.data(), rows, cols, device); + if (std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr) { + // Prime the persistent-zero cache for this spec during the eager warmup + // (capture-safe zeroing replays ttnn::copy(zero, dst) — see MemsetDevice). + ZeroCachePrime(ttnn::Shape({rows, cols}), ttnn::DataType::BFLOAT16, + ttnn::Layout::TILE, device); + } std::lock_guard g(SlotMutex()); BufferSlot* s = FindSlot(t.data); if (s != nullptr) { @@ -393,6 +499,8 @@ bool TryDevicePagedFill(ttnn::Tensor& cache_dev, MeshDevice& device, std::memcpy(dst, src, static_cast(d) * sizeof(float)); } } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryDevicePagedFill from_vector WRITE during capture\n"); ttnn::Tensor xt = ttnn::Tensor::from_vector( x, SpecOf(tt::tt_metal::Shape({1u, nkv, T_pad, d}), ttnn::DataType::BFLOAT16, ttnn::Layout::TILE), @@ -404,6 +512,8 @@ bool TryDevicePagedFill(ttnn::Tensor& cache_dev, MeshDevice& device, const uint32_t tok_i = std::min(j * bs, T - 1u); pt[static_cast(j)] = static_cast(blocks[static_cast(tok_i)]); } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryDevicePagedFill from_vector WRITE during capture\n"); ttnn::Tensor page_table = ttnn::Tensor::from_vector( pt, SpecOf(tt::tt_metal::Shape({1u, n_logical}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), @@ -431,6 +541,8 @@ ttnn::Tensor MakeHeightShardedUpdateInput(MeshDevice& device, const float* toks, static_cast(d) * sizeof(float)); } } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] MakeHeightShardedUpdateInput from_vector WRITE during capture\n"); ttnn::Tensor xt = ttnn::Tensor::from_vector( x, SpecOf(tt::tt_metal::Shape({1u, C, nkv_pad, d}), ttnn::DataType::BFLOAT16, ttnn::Layout::TILE), @@ -471,6 +583,8 @@ bool TryDevicePagedUpdateBatch(ttnn::Tensor& cache_dev, MeshDevice& device, pt[static_cast(b)] = static_cast(phys_blocks[static_cast(base + b)]); update_idxs[static_cast(b)] = offsets[static_cast(base + b)]; } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryDevicePagedUpdateBatch from_vector WRITE during capture\n"); ttnn::Tensor page_table = ttnn::Tensor::from_vector( pt, SpecOf(tt::tt_metal::Shape({C, 1u}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); @@ -513,6 +627,8 @@ bool TryDevicePagedFusedUpdateBatch(ttnn::Tensor& k_dev, ttnn::Tensor& v_dev, Me pt[static_cast(b)] = static_cast(phys_blocks[static_cast(base + b)]); update_idxs[static_cast(b)] = offsets[static_cast(base + b)]; } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryDevicePagedFusedUpdateBatch from_vector WRITE during capture\n"); ttnn::Tensor page_table = ttnn::Tensor::from_vector( pt, SpecOf(tt::tt_metal::Shape({C, 1u}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); @@ -698,6 +814,8 @@ ttnn::Tensor EnsurePagedKvTtnn(const Tensor& cache_nhd, MeshDevice& device, uint const auto spec = SpecOf(tt::tt_metal::Shape({used_nb, nkv, bs, d}), ttnn::DataType::BFLOAT16, ttnn::Layout::TILE); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] EnsurePagedKvTtnn from_vector WRITE during capture\n"); ttnn::Tensor dev = ttnn::Tensor::from_vector(upload, spec, &device); std::lock_guard g(PagedKvMutex()); @@ -758,6 +876,7 @@ void CommitHost(Tensor& out) { // Device compute: keep result on device (CommitDevice2D). Host round-trip only // when the consumer is a host-staged op (EnsureHost) or an untracked buffer. void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { + TT_OP_TRACE("Matmul"); VT_CHECK(a.rank == 2 && b.rank == 2 && out.rank == 2, "tenstorrent kMatmul: only rank-2 tensors are supported in W0"); VT_CHECK(IsFloatDType(a.dtype) && IsFloatDType(b.dtype) && @@ -784,6 +903,7 @@ void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { // that flag flipped — no separate upload shape needed since `b` is uploaded // in its native [N,K] layout and ttnn transposes on device. void MatmulBTKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { + TT_OP_TRACE("MatmulBT"); VT_CHECK(a.rank == 2 && b.rank == 2 && out.rank == 2, "tenstorrent kMatmulBT: only rank-2 tensors are supported in W0"); VT_CHECK(IsFloatDType(a.dtype) && IsFloatDType(b.dtype) && @@ -812,6 +932,7 @@ void MatmulBTKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { // this kernel's behavior pinned to the CPU reference rather than to // whatever ttnn::add happens to support today. void AddKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { + TT_OP_TRACE("Add"); VT_CHECK(a.rank == 2 && out.rank == 2, "tenstorrent kAdd: `a`/`out` must be rank-2 in W0"); VT_CHECK(b.rank == 2 || b.rank == 1, "tenstorrent kAdd: `b` must be rank-1 or rank-2 in W0"); VT_CHECK(IsFloatDType(a.dtype) && IsFloatDType(b.dtype) && @@ -835,6 +956,8 @@ void AddKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) { for (uint32_t r = 0; r < rows; ++r) for (uint32_t c = 0; c < d; ++c) replicated[static_cast(r) * d + c] = LoadElemF32(b, c); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] AddKernel from_vector WRITE during capture\n"); dev_b = ttnn::Tensor::from_vector(replicated, TileSpecOf(rows, d), &device); } else { dev_b = EnsureDevice2D(b, device); @@ -895,6 +1018,8 @@ ttnn::Tensor EnsureEmbedTableDevice(const Tensor& table, MeshDevice& device) { } EnsureHost(table); std::vector host_table = ToHostF32(table); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] EnsureEmbedTableDevice from_vector WRITE during capture\n"); ttnn::Tensor dev_table = ttnn::Tensor::from_vector( host_table, SpecOf(tt::tt_metal::Shape({vocab, h}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), @@ -917,6 +1042,7 @@ ttnn::Tensor EnsureEmbedTableDevice(const Tensor& table, MeshDevice& device) { // vt::EmbeddingFn's (table, ids). Output is TILE so the next matmul can keep // the activation device-resident without a host round-trip. void EmbeddingKernel(Queue&, Tensor& out, const Tensor& table, const Tensor& ids) { + TT_OP_TRACE("Embedding"); VT_CHECK(table.rank == 2 && ids.rank == 1 && out.rank == 2, "tenstorrent kEmbedding: table rank-2, ids rank-1, out rank-2"); VT_CHECK(IsFloatDType(table.dtype) && (out.dtype == DType::kF32 || out.dtype == DType::kBF16), @@ -948,6 +1074,8 @@ void EmbeddingKernel(Queue&, Tensor& out, const Tensor& table, const Tensor& ids } } MeshDevice& device = SharedMeshDevice(); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] EmbeddingKernel from_vector WRITE during capture\n"); ttnn::Tensor dev_ids = ttnn::Tensor::from_vector( host_ids, SpecOf(tt::tt_metal::Shape({t}), ttnn::DataType::UINT32, ttnn::Layout::ROW_MAJOR), &device); @@ -981,6 +1109,8 @@ ttnn::Tensor EnsureAffine1D(const Tensor& t, uint32_t d, MeshDevice& device) { EnsureHost(t); std::vector host(d); for (uint32_t i = 0; i < d; ++i) host[i] = LoadElemF32(t, i); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] EnsureAffine1D from_vector WRITE during capture\n"); ttnn::Tensor dev = ttnn::Tensor::from_vector( host, SpecOf(tt::tt_metal::Shape({1, d}), ttnn::DataType::BFLOAT16, ttnn::Layout::TILE), &device); @@ -1004,6 +1134,7 @@ ttnn::Tensor EnsureAffine1D(const Tensor& t, uint32_t d, MeshDevice& device) { // LayerNormArgs (OPT default 1e-5, not ttnn's 1e-12 default). void LayerNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor* weight, const Tensor* bias, const LayerNormArgs& args) { + TT_OP_TRACE("LayerNorm"); VT_CHECK(x.rank == 2 && out.rank == 2, "tenstorrent kLayerNorm: only rank-2 tensors are supported in this step"); VT_CHECK(IsFloatDType(x.dtype) && (out.dtype == DType::kF32 || out.dtype == DType::kBF16), @@ -1044,6 +1175,7 @@ void LayerNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor* weight, // upload via the same TILE [1,D] affine helper as kLayerNorm. void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& weight, const RmsNormArgs& args, Tensor* residual) { + TT_OP_TRACE("RmsNorm"); VT_CHECK(x.rank == 2 && out.rank == 2, "tenstorrent kRmsNorm: only rank-2 tensors are supported in this step"); VT_CHECK(IsFloatDType(x.dtype) && (out.dtype == DType::kF32 || out.dtype == DType::kBF16), @@ -1068,8 +1200,14 @@ void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& weight, // (rows=1) pays more for device add+rms launches than a host loop, and was // a measurable e2e regression vs host residual. constexpr uint32_t kDeviceResidualMinRows = 32; - const bool host_residual = - args.gemma || (residual != nullptr && rows < kDeviceResidualMinRows); + // HOST-FREE-FORWARD R1: force the residual merge + RMS device path at T=1 when + // capture is desired (ttnn trace prohibits host ops in the captured region). + // Opt-in via VT_TT_HOST_FREE_DECODE; inert by default (keeps the 12.5 tok/s + // hybrid baseline). Numerics proven by BACKEND-TENSTORRENT-RESIDUAL-GOLDEN. + const bool host_free_decode = + std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr; + const bool host_residual = !host_free_decode && + (args.gemma || (residual != nullptr && rows < kDeviceResidualMinRows)); if (host_residual) { EnsureHost(x); EnsureHost(weight); @@ -1124,6 +1262,7 @@ void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& weight, // last-dim halves, ttnn::silu(gate), ttnn::multiply by up. BF16 tile path // (same envelope as matmul/norm); not bit-exact vs host f32. void SiluAndMulKernel(Queue&, Tensor& out, const Tensor& x) { + TT_OP_TRACE("SiluAndMul"); VT_CHECK(x.rank == 2 && out.rank == 2, "tenstorrent kSiluAndMul: only rank-2 tensors are supported"); VT_CHECK(IsFloatDType(x.dtype) && (out.dtype == DType::kF32 || out.dtype == DType::kBF16), @@ -1156,6 +1295,7 @@ void SiluAndMulKernel(Queue&, Tensor& out, const Tensor& x) { // CastBf16Kernel / CastF32Kernel). Qwen3 uses these for K/V cache dtype and // the logits / rope-cache paths. Host-staged; bit-exact for supported pairs. void CastBf16Kernel(Queue&, Tensor& out, const Tensor& in) { + TT_OP_TRACE("CastBf16"); VT_CHECK(out.dtype == DType::kBF16, "tenstorrent kCastBf16: out must be bf16"); VT_CHECK(IsFloatDType(in.dtype), "tenstorrent kCastBf16: in must be float"); VT_CHECK(out.Numel() == in.Numel(), "tenstorrent kCastBf16: numel mismatch"); @@ -1168,6 +1308,7 @@ void CastBf16Kernel(Queue&, Tensor& out, const Tensor& in) { } void CastF32Kernel(Queue&, Tensor& out, const Tensor& in) { + TT_OP_TRACE("CastF32"); VT_CHECK(out.dtype == DType::kF32, "tenstorrent kCastF32: out must be f32"); VT_CHECK(IsFloatDType(in.dtype), "tenstorrent kCastF32: in must be float"); VT_CHECK(out.Numel() == in.Numel(), "tenstorrent kCastF32: numel mismatch"); @@ -1236,8 +1377,46 @@ void RopeApplyDeviceNeox(Tensor& x3, const float* cos_t, const float* sin_t, const uint32_t halfu = static_cast(half); const uint32_t rotu = static_cast(rot); const uint32_t du = static_cast(d); - ttnn::Tensor dev_cos = UploadRows(cos_exp.data(), thu, halfu, device); - ttnn::Tensor dev_sin = UploadRows(sin_exp.data(), thu, halfu, device); + // ITEM 5: persistent cos/sin — build outside capture, copy in-region. + const std::string rk = RopeCSKey(thu, halfu); + ttnn::Tensor dev_cos, dev_sin; + bool cache_hit = false; + { + std::lock_guard g(RopeCSMutex()); + auto it = RopeCSCache().find(rk); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, + "[TT-TRACE] rope lookup key=%s found=%d content_eq=%d " + "(want first=%f n=%zu)\n", + rk.c_str(), it != RopeCSCache().end(), + it != RopeCSCache().end() && it->second.cos_host == cos_exp, + cos_exp.empty() ? -1.0f : cos_exp.front(), cos_exp.size()); + if (it != RopeCSCache().end() && it->second.cos_host == cos_exp) { + dev_cos = it->second.cos; + dev_sin = it->second.sin; + cache_hit = true; + } + } + if (!cache_hit) { + VT_CHECK(!tt_capture_active(), + "tenstorrent: rope cos/sin cache miss during capture — the " + "table changed (positions moved); the decode-graph driver must " + "call WarmRopeCosSin for the step's positions BEFORE BeginCapture " + "(the SizeSlot::Refresh pattern)"); + dev_cos = UploadRows(cos_exp.data(), thu, halfu, device); + dev_sin = UploadRows(sin_exp.data(), thu, halfu, device); + std::lock_guard g(RopeCSMutex()); + RopeCSEntry e; + e.cos = dev_cos; + e.sin = dev_sin; + e.cos_host = cos_exp; + RopeCSCache()[rk] = std::move(e); + } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] rope cos/sin cache %s during capture " + "(key th=%u half=%u first=%f)\n", + cache_hit ? "HIT" : "MISS", thu, halfu, + cos_exp.empty() ? -1.0f : cos_exp.front()); // x1 = x[..., :half], x2h = x[..., half:rot] (NeoX half-split) ttnn::Tensor x1 = ttnn::slice(dev_x, ttsl::SmallVector{0, 0}, @@ -1340,11 +1519,14 @@ void RopeApplyHost(Tensor& qs, Tensor* ks, const float* cos_t, const float* sin_ // Short Qwen3 decode (T=1,H=16) is host-faster even when Q is already on device // (measured regression when always-device-for-resident was forced). inline bool PreferDeviceRope(int64_t tokens, int64_t heads) { + // HOST-FREE-FORWARD R1: force device RoPE at T=1 for capture (see RmsNorm note). + if (std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr) return true; return tokens * heads >= 64; } // kRopeNeox: Qwen3-dense RoPE. Device NeoX for large [T*H]; host for short decode. void RopeNeoxKernel(Queue&, Tensor& qs, Tensor& ks, const Tensor& pos, const RopeArgs& args) { + TT_OP_TRACE("RopeNeox"); VT_CHECK(qs.rank == 3 && ks.rank == 3, "tenstorrent kRopeNeox: qs/ks rank-3"); VT_CHECK(IsFloatDType(qs.dtype) && qs.dtype == ks.dtype, "tenstorrent kRopeNeox: qs/ks float same dtype"); @@ -1362,6 +1544,12 @@ void RopeNeoxKernel(Queue&, Tensor& qs, Tensor& ks, const Tensor& pos, const Rop std::vector cos_t, sin_t; BuildCosSinFromPositions(pos, t, args.rotary_dim, static_cast(args.base), args, cos_t, sin_t); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] rope kernel pos0=%d t=%lld hq=%lld cos_first=%f\n", + (int)(pos.dtype == DType::kI32 ? pos.Ptr()[0] + : static_cast(pos.Ptr()[0])), + (long long)t, (long long)hq, + cos_t.empty() ? -1.0f : cos_t.front()); if (PreferDeviceRope(t, hq)) { MeshDevice& device = SharedMeshDevice(); RopeApplyDeviceNeox(qs, cos_t.data(), sin_t.data(), t, hq, d, args.rotary_dim, device); @@ -1458,6 +1646,7 @@ void RopeFromCacheKernel(Queue&, Tensor& qs, Tensor* ks, const Tensor& positions // reuse without download+reupload. Host path (bit-exact memcpy) when qkv is // host-only — unit tests and weight-load style callers. void QkvSplitKernel(Queue&, Tensor& q_out, Tensor& k_out, Tensor& v_out, const Tensor& qkv) { + TT_OP_TRACE("QkvSplit"); VT_CHECK(qkv.rank == 2 && IsFloatDType(qkv.dtype), "tenstorrent kQkvSplit: rank-2 float qkv required"); VT_CHECK(q_out.dtype == qkv.dtype && k_out.dtype == qkv.dtype && v_out.dtype == qkv.dtype, @@ -1524,20 +1713,273 @@ void QkvSplitKernel(Queue&, Tensor& q_out, Tensor& k_out, Tensor& v_out, const T // (cpu_cache.cpp ReshapeAndCacheKernel). Stride-driven so unbind-style // [num_blocks,2,bs,H,D] views work; slot < 0 is a padded-token skip. // Host-staged pure element copy for F32. +// ---- ITEM 5 (PA): persistent page_table + cur_pos device tensors ------------- +namespace { +struct PaMetaEntry { + ttnn::Tensor page_table; // int32 [B, max_blocks] device + ttnn::Tensor cur_pos; // int32 [B] device + std::vector pt_host; + std::vector cp_host; +}; +std::mutex& PaMetaMutex() { static std::mutex m; return m; } +std::map, PaMetaEntry>& PaMetaCache() { + static std::map, PaMetaEntry> c; + return c; +} +} // namespace + +// ---- ITEM 5 (RAC): persistent update-idx / page-table device tensors ------- +// Refreshed by WarmRacIdx (driver Refresh slot, outside capture) so the +// captured paged_update_cache replays against stable addresses. Keyed by the +// slot-mapping HOST buffer (the decode-graph slot's persistent buffer), so a +// different graph size gets its own entries. +namespace { +struct RacIdxEntry { + ttnn::Tensor update_idxs; // int32 [C] device + ttnn::Tensor page_table; // int32 [C,1] device + std::vector idx_host; // content identity for reuse check + ttnn::Tensor sharded_zero; // pre-allocated height-sharded [1,1,nkv_pad,d] zero +}; +std::mutex& RacIdxMutex() { static std::mutex m; return m; } +// Keyed by (num_slots, block_size) shape — idx tensors depend on slot values + block_size. +// Pending RAC: when capture skips the KV write, record the k/v device +// shadows + slot so the driver can flush them at the next Refresh (outside +// capture), then replay uses the updated KV. +struct PendingRac { + ttnn::Tensor k_dev, v_dev; // device shadows of the rope output + void* k_cache = nullptr; + void* v_cache = nullptr; +}; +std::mutex& PendingRacMutex() { static std::mutex m; return m; } +// One entry per LAYER (each layer has its own attn_kv[l].data cache pair). +std::vector& PendingRacVec() { static std::vector v; return v; } +int64_t& PendingRacSlotId() { static int64_t s = -1; return s; } + +std::map, RacIdxEntry>& RacIdxCache() { + static std::map, RacIdxEntry> c; + return c; +} +} // namespace + +// Warm hook: stage persistent idx tensors for THIS slot mapping. Host reads +// here are legal (called outside capture). Idempotent per content change. + +// Host-free decode RAC: device shadows in, paged_update_cache out. Returns +// false (host path) unless every precondition holds. +bool TryReshapeAndCacheDeviceDecode(const Tensor& k, const Tensor& v, + Tensor& k_cache, Tensor& v_cache, + const Tensor& slot_mapping) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] TryRACDevice called capturing=%d\n", + (int)tt_capture_active()); + const int64_t T = k.shape[0]; + const int64_t nkv = k.shape[1]; + const int64_t d = k.shape[2]; + const int64_t bs = k_cache.shape[1]; + const int64_t num_slots = slot_mapping.shape[0]; + if (T < 1 || num_slots < 1) return false; + if ((d % 32u) != 0u || (bs % 32u) != 0u) return false; + if (num_slots > 1) return false; // decode T=1 only for now + + // k/v must carry CURRENT device shadows ([T*nkv, d] TILE bf16 from rope). + std::optional k_dev, v_dev; + { + std::lock_guard g(SlotMutex()); + BufferSlot* sk = FindSlot(k.data); + BufferSlot* sv = FindSlot(v.data); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] RAC kv shadow k_ptr=%p sk=%p dc=%d val=%d | v_ptr=%p sv=%p dc=%d val=%d\n", + k.data, (void*)sk, sk?sk->device_current:0, sk?(int)sk->device.has_value():0, + v.data, (void*)sv, sv?sv->device_current:0, sv?(int)sv->device.has_value():0); + if (sk == nullptr || !sk->device_current || !sk->device.has_value()) return false; + if (sv == nullptr || !sv->device_current || !sv->device.has_value()) return false; + k_dev = sk->device; + v_dev = sv->device; + } + + // Paged-KV shadows must exist and cover the target block. + const int64_t slot = slot_mapping.Ptr()[0]; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] RAC slot=%lld cap=%d\n", + (long long)slot, (int)tt_capture_active()); + if (slot < 0) return true; // nothing to write; treat as handled + const uint32_t block = static_cast(slot / bs); + const uint32_t offset = static_cast(slot % bs); + + std::optional kc_dev, vc_dev; + { + std::lock_guard g(PagedKvMutex()); + PagedKvShadow* skc = &PagedKvShadows()[reinterpret_cast(k_cache.data)]; + PagedKvShadow* svc = &PagedKvShadows()[reinterpret_cast(v_cache.data)]; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] RAC paged-kv shadow k=%d v=%d k_nb=%u\n", + skc->device.has_value(), svc->device.has_value(), skc->nb); + if (!skc->device.has_value() || !svc->device.has_value()) return false; + if (skc->nb <= block || skc->nkv != static_cast(nkv) || + skc->bs != static_cast(bs) || skc->d != static_cast(d)) return false; + if (svc->nb <= block || svc->nkv != static_cast(nkv) || + svc->bs != static_cast(bs) || svc->d != static_cast(d)) return false; + kc_dev = skc->device; + vc_dev = svc->device; + } + + // Persistent idx tensors for THIS slot-mapping buffer (warmed outside + // capture). Both must exist; content refresh happens at warm time. + ttnn::Tensor idx, pt; + { + std::lock_guard g(RacIdxMutex()); + const auto key = std::make_pair(num_slots, static_cast(bs)); + auto it = RacIdxCache().find(key); + const int64_t slot0 = slot_mapping.Ptr()[0]; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] RAC idx-check slot0=%lld cap=%d key=(%lld,%lld)\n", + (long long)slot0, (int)tt_capture_active(), + (long long)num_slots, (long long)bs); + const int32_t expect_idx = (slot0 < 0) ? 0 : static_cast(slot0 % bs); + if (it == RacIdxCache().end() || + (it->second.idx_host.size() >= 1 && it->second.idx_host[0] != expect_idx)) { + VT_CHECK(!tt_capture_active(), + "tenstorrent: RAC idx tensors not warmed — call WarmRacIdx " + "outside capture (driver Refresh slot) first"); + return false; + } + idx = it->second.update_idxs; + pt = it->second.page_table; + } + + // During capture: skip the KV write (to_memory_config triggers the DRAM + // overlap fatal). The driver flushes the pending k/v at the Refresh slot. + // One-step lag (PA reads previous step's KV), but capture+replay works. + if (tt_capture_active()) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] RAC skip during capture (flush at refresh)\n"); + std::lock_guard g(PendingRacMutex()); + PendingRac e; + e.k_dev = *k_dev; + e.v_dev = *v_dev; + e.k_cache = k_cache.data; + e.v_cache = v_cache.data; + PendingRacVec().push_back(std::move(e)); + PendingRacSlotId() = slot; + return true; + } + + // Eager (cold) step: do the real device RAC here (safe outside capture). + MeshDevice& device = SharedMeshDevice(); + const uint32_t nkv_pad = std::max(32u, ((static_cast(nkv) + 31u) / 32u) * 32u); + + const RacIdxEntry& rac_entry = [&] { + std::lock_guard g(RacIdxMutex()); + return RacIdxCache().at(std::make_pair(num_slots, static_cast(bs))); + }(); + + auto build_input = [&](const ttnn::Tensor& src) -> ttnn::Tensor { + // src: [T*nkv, d] TILE bf16 (T==1 -> [nkv, d]). + // During capture: copy the rope output into the PRE-BUILT sharded_zero + // (ttnn::copy with the sharded output; paged_update_cache is in-place). + // During eager: build freely (zeros+concat+shard are safe). + if (tt_capture_active()) { + // Copy src (TILE) into the sharded_zero's top rows via ttnn::copy + // with an explicit output memory config. May allocate if layout + // conversion is needed — but paged_update_cache is in-place, so the + // overall effect is correct even if the copy allocates a temp. + ttnn::Tensor reshaped = src.reshape(ttnn::Shape( + {1u, 1u, static_cast(nkv), static_cast(d)})); + const uint32_t pad_heads = nkv_pad - static_cast(nkv); + ttnn::Tensor padded; + if (pad_heads > 0) { + ttnn::Tensor zeros_tail = ZeroCacheGet( + ttnn::zeros(ttnn::Shape({1u, 1u, pad_heads, static_cast(d)}), + ttnn::DataType::BFLOAT16, ttnn::Layout::TILE, std::ref(device)), + device); + padded = ttnn::concat(std::vector{reshaped, zeros_tail}, 2); + } else { + padded = reshaped; + } + // Copy the padded TILE tensor into the PRE-BUILT sharded buffer. + // ttnn::copy with an existing sharded destination should write + // in-place (no allocation). This is the capture-safe path. + ttnn::copy(padded, rac_entry.sharded_zero); + return rac_entry.sharded_zero; + } + // Eager: build freely. + ttnn::Tensor reshaped = src.reshape(ttnn::Shape( + {1u, 1u, static_cast(nkv), static_cast(d)})); + const uint32_t pad_heads = nkv_pad - static_cast(nkv); + ttnn::Tensor padded; + if (pad_heads > 0) { + ttnn::Tensor zeros_tail = ttnn::zeros( + ttnn::Shape({1u, 1u, pad_heads, static_cast(d)}), + ttnn::DataType::BFLOAT16, ttnn::Layout::TILE, std::ref(device)); + padded = ttnn::concat(std::vector{reshaped, zeros_tail}, 2); + } else { + padded = reshaped; + } + const auto grid = device.compute_with_storage_grid_size(); + const tt::tt_metal::CoreRangeSet core_set = + tt::tt_metal::num_cores_to_corerangeset(1u, grid, true); + tt::tt_metal::ShardSpec ss(core_set, {nkv_pad, static_cast(d)}, + tt::tt_metal::ShardOrientation::ROW_MAJOR); + tt::tt_metal::MemoryConfig sm(tt::tt_metal::TensorMemoryLayout::HEIGHT_SHARDED, + tt::tt_metal::BufferType::L1, ss); + return ttnn::to_memory_config(padded, sm); + }; + + ttnn::Tensor k_in = build_input(*k_dev); + ttnn::Tensor v_in = build_input(*v_dev); + ttnn::Tensor new_kc = ttnn::experimental::paged_update_cache( + *kc_dev, k_in, /*update_idxs=*/{}, rac_entry.update_idxs, + /*share_cache=*/false, rac_entry.page_table, + /*batch_offset=*/0, /*compute_kernel_config=*/std::nullopt, + /*mesh_coords=*/std::nullopt); + ttnn::Tensor new_vc = ttnn::experimental::paged_update_cache( + *vc_dev, v_in, /*update_idxs=*/{}, rac_entry.update_idxs, + /*share_cache=*/false, rac_entry.page_table, + /*batch_offset=*/0, /*compute_kernel_config=*/std::nullopt, + /*mesh_coords=*/std::nullopt); + { + std::lock_guard g(PagedKvMutex()); + PagedKvShadows()[reinterpret_cast(k_cache.data)].device = std::move(new_kc); + PagedKvShadows()[reinterpret_cast(k_cache.data)].device_current = true; + PagedKvShadows()[reinterpret_cast(v_cache.data)].device = std::move(new_vc); + PagedKvShadows()[reinterpret_cast(v_cache.data)].device_current = true; + } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] RAC eager device update (cold step)\n"); + (void)offset; (void)block; (void)nkv_pad; + return true; +} + void ReshapeAndCacheKernel(Queue&, const Tensor& k, const Tensor& v, Tensor& k_cache, Tensor& v_cache, const Tensor& slot_mapping) { + TT_OP_TRACE("ReshapeAndCache"); VT_CHECK(k.rank == 3 && v.rank == 3 && k_cache.rank == 4 && v_cache.rank == 4, "tenstorrent kReshapeAndCache: k/v rank-3, caches rank-4"); VT_CHECK(IsFloatDType(k.dtype) && k.dtype == v.dtype && k_cache.dtype == k.dtype && v_cache.dtype == k.dtype, "tenstorrent kReshapeAndCache: k/v/caches must share one float dtype"); + VT_CHECK(slot_mapping.rank == 1 && slot_mapping.dtype == DType::kI64, + "tenstorrent kReshapeAndCache: slot_mapping rank-1 i64"); + + // ITEM 5 (RAC): host-free decode branch. The host path below downloads k/v + // (rope output shadows) and re-uploads via from_vector in the device push — + // both fatal during capture. This branch instead feeds the DEVICE shadows + // straight into paged_update_cache with persistent idx/page-table tensors. + // Conditions: capturing (or host-free flag), all inputs device-shadowed, + // TILE-legal dims, and the warm hook already staged the idx tensors. + static const bool host_free_rac = + std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr; + if (host_free_rac || tt_capture_active()) { + if (TryReshapeAndCacheDeviceDecode(k, v, k_cache, v_cache, slot_mapping)) { + return; + } + } + EnsureHost(k); EnsureHost(v); EnsureHost(k_cache); EnsureHost(v_cache); EnsureHost(slot_mapping); - VT_CHECK(slot_mapping.rank == 1 && slot_mapping.dtype == DType::kI64, - "tenstorrent kReshapeAndCache: slot_mapping rank-1 i64"); const int64_t num_slots = slot_mapping.shape[0]; const int64_t block_size = k_cache.shape[1]; const int64_t num_kv_heads = k_cache.shape[2]; @@ -1623,11 +2065,14 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso const Tensor& v_cache, const Tensor& block_table, const Tensor& seq_lens, const Tensor& query_start_loc, const PagedAttentionArgs& args) { + TT_OP_TRACE("TryPagedAttentionDeviceDecode"); if (!args.causal || args.logits_soft_cap > 0.0f) return false; if (args.window_size.has_value()) return false; if (args.kv_cache_dtype != Fp8KVCacheDataType::kAuto) return false; if (query.rank != 3 || out.rank != 3 || k_cache.rank != 4 || v_cache.rank != 4) return false; if (!query.IsContiguous() || !out.IsContiguous()) return false; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] TryPADecode entered cap=%d\n", (int)tt_capture_active()); const int64_t total_q = query.shape[0]; const int64_t hq = query.shape[1]; @@ -1674,9 +2119,40 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso if (static_cast(used_nb) > k_cache.shape[0]) return false; try { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA reached EnsurePagedKvTtnn cap=%d used_nb=%u\n", (int)tt_capture_active(), used_nb); + MeshDevice& device = SharedMeshDevice(); - ttnn::Tensor dev_k = EnsurePagedKvTtnn(k_cache, device, used_nb); - ttnn::Tensor dev_v = EnsurePagedKvTtnn(v_cache, device, used_nb); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] PA EnsurePagedKvTtnn k used_nb=%u\n", used_nb); + // Use the cached shadow when it exists (primed by WarmPagedKvShadow). + // This skips EnsurePagedKvTtnn's from_vector upload AND its contiguous + // check (KvSlice returns a non-contiguous strided view that the VT_CHECK + // rejects). Needed on BOTH cold and capture steps so sdpa_decode compiles. + ttnn::Tensor dev_k, dev_v; + { + std::lock_guard g(PagedKvMutex()); + auto& sk = PagedKvShadows()[reinterpret_cast(k_cache.data)]; + auto& sv = PagedKvShadows()[reinterpret_cast(v_cache.data)]; + if (sk.device_current && sk.device.has_value() && sk.nb >= used_nb && + sv.device_current && sv.device.has_value() && sv.nb >= used_nb) { + dev_k = *sk.device; + dev_v = *sv.device; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA using cached KV shadows (k_nb=%u v_nb=%u) cap=%d\n", + sk.nb, sv.nb, (int)tt_capture_active()); + } else if (tt_capture_active()) { + throw std::runtime_error("PA: no KV shadow during capture"); + } else { + // Cold step without shadow: fall through to EnsurePagedKvTtnn + // (may fail on non-contiguous KvSlice; that's OK — the host path runs). + g.~lock_guard(); // release before EnsurePagedKvTtnn + dev_k = EnsurePagedKvTtnn(k_cache, device, used_nb); + dev_v = EnsurePagedKvTtnn(v_cache, device, used_nb); + } + } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] PA KV shadows OK, building page_table\n"); const uint32_t Bu = static_cast(num_reqs); const uint32_t hu = static_cast(hq); @@ -1684,6 +2160,10 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso // Q: [1, B, H, D]. Prefer reshape of a resident [B*H, D] / [B, H*D] shadow // (post device rope) so we never download then re-upload. + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA before identity_q cap=%d total_q=%lld num_reqs=%lld qsl0=%d qsl1=%d\n", + (int)tt_capture_active(), (long long)total_q, (long long)num_reqs, + qsl[0], num_reqs > 0 ? qsl[1] : -1); bool identity_q = true; for (int64_t r = 0; r < num_reqs; ++r) { if (qsl[r] != r) { @@ -1699,7 +2179,11 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso ttnn::Tensor dev_q2d = EnsureDevice2D(q_flat, device); dev_q = ttnn::reshape(dev_q2d, ttnn::Shape({1u, Bu, hu, du})); q_from_device = true; - } catch (const std::exception&) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA q_from_device OK cap=%d\n", (int)tt_capture_active()); + } catch (const std::exception& e) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA q_from_device FAILED: %s\n", e.what()); q_from_device = false; } } @@ -1715,22 +2199,43 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso } } } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryPagedAttentionDeviceDecode from_vector WRITE during capture\n"); dev_q = ttnn::Tensor::from_vector( q_host, SpecOf(tt::tt_metal::Shape({1u, Bu, hu, du}), ttnn::DataType::BFLOAT16, ttnn::Layout::TILE), &device); } - ttnn::Tensor dev_pt = ttnn::Tensor::from_vector( - pt, SpecOf(tt::tt_metal::Shape({Bu, static_cast(max_blocks)}), - ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), - &device); - - // cur_pos [B] = seq_len - 1 - std::vector cpos(static_cast(num_reqs)); - for (int64_t r = 0; r < num_reqs; ++r) cpos[static_cast(r)] = slens[r] - 1; - ttnn::Tensor dev_pos = ttnn::Tensor::from_vector( - cpos, SpecOf(tt::tt_metal::Shape({Bu}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), - &device); + ttnn::Tensor dev_pt, dev_pos; + if (tt_capture_active()) { + // Use persistent PA metadata (warmed at the Refresh slot). + std::lock_guard g(PaMetaMutex()); + const auto pkey = std::make_pair(static_cast(num_reqs), + static_cast(max_blocks)); + auto it = PaMetaCache().find(pkey); + const int32_t expect_cp = slens[0] - 1; + if (it == PaMetaCache().end() || + (it->second.cp_host.size() >= 1 && it->second.cp_host[0] != expect_cp)) { + VT_CHECK(false, "tenstorrent: PA meta not warmed for this step"); + } + dev_pt = it->second.page_table; + dev_pos = it->second.cur_pos; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA using cached meta (pt+cp)\n"); + } else { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryPagedAttentionDeviceDecode from_vector WRITE during capture\n"); + dev_pt = ttnn::Tensor::from_vector( + pt, SpecOf(tt::tt_metal::Shape({Bu, static_cast(max_blocks)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + // cur_pos [B] = seq_len - 1 + std::vector cpos(static_cast(num_reqs)); + for (int64_t r = 0; r < num_reqs; ++r) cpos[static_cast(r)] = slens[r] - 1; + dev_pos = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({Bu}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + } const auto grid = device.compute_with_storage_grid_size(); ttnn::operations::transformer::SDPAProgramConfig prog{ @@ -1741,6 +2246,8 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso /*exp_approx_mode=*/false, /*max_cores_per_head_batch=*/16}; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA calling sdpa_decode cap=%d\n", (int)tt_capture_active()); ttnn::Tensor dev_out = ttnn::transformer::paged_scaled_dot_product_attention_decode( dev_q, dev_k, dev_v, dev_pt, /*is_causal=*/true, @@ -1778,6 +2285,8 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso // Output ~ [1, B, H, D] → host [B, H, D] in request order, then scatter to // global query token indices. + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA decode to_vector\n"); std::vector result = dev_out.to_vector(); VT_CHECK(static_cast(result.size()) >= num_reqs * hq * d, "tenstorrent device PA: unexpected output size"); @@ -1792,7 +2301,9 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso } CommitHost(out); return true; - } catch (const std::exception&) { + } catch (const std::exception& e) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA device decode FAILED: %s\n", e.what()); // Fall back to host oracle (shape/grid/dtype edge cases). return false; } @@ -1908,6 +2419,8 @@ bool TryPagedAttentionDevicePrefill(Tensor& out, const Tensor& query, const Tens const int64_t need_kv = chunk_start0 + q_pad; if (max_blocks * block_size < need_kv) return false; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryPagedAttentionDevicePrefill from_vector WRITE during capture\n"); ttnn::Tensor dev_pt = ttnn::Tensor::from_vector( pt, SpecOf(tt::tt_metal::Shape({1u, static_cast(max_blocks)}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), @@ -1927,6 +2440,8 @@ bool TryPagedAttentionDevicePrefill(Tensor& out, const Tensor& query, const Tens } } } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-UP] TryPagedAttentionDevicePrefill from_vector WRITE during capture\n"); ttnn::Tensor dev_q = ttnn::Tensor::from_vector( q_host, SpecOf(tt::tt_metal::Shape({1u, hu, static_cast(kChunk), du}), @@ -1954,6 +2469,8 @@ bool TryPagedAttentionDevicePrefill(Tensor& out, const Tensor& query, const Tens continue; } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA prefill to_vector\n"); std::vector result = dev_out.to_vector(); // Expected dense logical [1, H, kChunk, D]. VT_CHECK(static_cast(result.size()) >= hq * kChunk * d, @@ -2372,6 +2889,9 @@ void TraceBeginCapture() { MeshDevice& device = SharedMeshDevice(); s.capturing_id = ttnn::operations::trace::begin_trace_capture(&device, kTraceCq); s.capturing = true; + tt_capture_active() = true; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] BeginCapture (flag set)\n"); } void TraceEndCapture() { @@ -2389,6 +2909,7 @@ void TraceEndCapture() { s.replay_id = s.capturing_id; s.has_replay = true; s.capturing = false; + tt_capture_active() = false; } void TraceReplay() { @@ -2405,6 +2926,7 @@ void* TraceEndCaptureGraph() { MeshDevice& device = SharedMeshDevice(); ttnn::operations::trace::end_trace_capture(&device, s.capturing_id, kTraceCq); s.capturing = false; + tt_capture_active() = false; // Opaque handle: heap-allocated MeshTraceId for the multi-graph API. return new ttnn::MeshTraceId(s.capturing_id); } @@ -2467,6 +2989,8 @@ void MarkHostWritten(void* host) { void EnsureHostBytes(void* host) { if (host == nullptr) return; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && tt_capture_active()) + std::fprintf(stderr, "[TT-TRACE] EnsureHostBytes DURING CAPTURE\n"); ttnn::Tensor dev; size_t bytes = 0; void* base = nullptr; @@ -2500,4 +3024,394 @@ void EnsureHostBytes(void* host) { } } +// ITEM 5: persistent zero tensors, created OUTSIDE capture (ttnn::zeros +// host-fills + to_device()s = an enqueue_write, illegal during trace capture). +// EnsureDevice2D primes the cache during the eager warmup so the captured +// res.Zero finds its entry and replays a warm device->device ttnn::copy. + +// HOST-FREE-FORWARD R2: device->device copy when capturing, so Backend::Copy +// does not to_vector inside the captured region. Both dst and src must carry a +// current device shadow of equal byte size; dst's shadow becomes a copy of src. +bool CopyDeviceDeviceIfCapture(void* dst, const void* src) { + // Run the device->device copy when EITHER capturing OR in host-free-decode + // mode (the env opt-in). The latter is essential so the EAGER warmup step + // (which the decode-graph framework runs BEFORE capture) also exercises + // ttnn::empty+ttnn::copy, compiling those programs into the cache so the + // subsequent capture doesn't hit "Cannot load new binaries during trace + // capture." + static const bool host_free = + std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr; + if (!tt_capture_active() && !host_free) return false; + static bool once = [&] { + // Enable program cache once on the first host-free path use — ttnn trace + // requires every captured op to be program-cache-warm. + MeshDevice& device = SharedMeshDevice(); + device.enable_program_cache(); + return true; + }(); + (void)once; + ttnn::Tensor src_dev; + { + std::lock_guard g(SlotMutex()); + BufferSlot* s = FindSlot(const_cast(src)); + BufferSlot* d = FindSlot(dst); + if (s == nullptr || !s->device_current || !s->device.has_value()) return false; + if (d == nullptr) return false; + if (s->bytes != d->bytes) return false; + src_dev = *s->device; + } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] device->device copy (capture-safe)\n"); + MeshDevice& device = SharedMeshDevice(); + // Allocate a destination device tensor matching src's shape/dtype/layout, + // then copy. No host readback. + ttnn::Tensor cloned = ttnn::empty(src_dev.logical_shape(), src_dev.dtype(), + src_dev.layout(), &device, + src_dev.memory_config()); + cloned = ttnn::copy(src_dev, cloned); + { + std::lock_guard g(SlotMutex()); + BufferSlot* d = FindSlot(dst); + if (d == nullptr) return false; + d->device = std::move(cloned); + d->device_current = true; + d->host_current = false; + } + return true; +} + +// HOST-FREE-FORWARD R3: on-device fill (for DBuf::Zero -> Backend::Memset) +// when host-free decode is active, so no host write happens inside capture. +// Reinterprets the buffer as a 2D [rows, cols] f32 tensor matching the +// existing device shadow's numel (zeros is the only value the forward uses). +bool MemsetDeviceIfCapture(void* p, int value) { + static const bool host_free = + std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr; + if (!tt_capture_active() && !host_free) return false; + if (value != 0) return false; // only zero-fill is handled on-device + // Need an existing shadow to know shape/dtype; or allocate from the slot. + std::optional dev; + { + std::lock_guard g(SlotMutex()); + BufferSlot* s = FindSlot(p); + if (s != nullptr && s->device_current && s->device.has_value()) { + dev = *s->device; + } + } + if (!dev.has_value()) { + // No shadow yet: DBuf::Zero on a brand-new buffer with no device tensor. + return false; // fall back to host memset; the buffer is host-only for now + } + MeshDevice& device = SharedMeshDevice(); + const ttnn::Tensor& shadow = *dev; + // ITEM 5: ttnn::zeros/full is NOT capture-safe — full_impl host-fills and + // to_device()s (creation.cpp:52-71), i.e. an enqueue_write that ttnn trace + // fatals on. The plugin pattern instead: keep PERSISTENT zero tensors + // (created outside capture, at warmup) and ttnn::copy one onto the target — + // a device->device program that is captured/replayed like any other warm op. + ttnn::Tensor zero_src = ZeroCacheGet(shadow, device); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] device zero-fill (capture-safe)\n"); + // Copy the persistent zero onto the shadow IN PLACE (keeps the shadow's + // device address stable — the whole point of persistent buffers). + ttnn::Tensor z = ttnn::copy(zero_src, shadow); + (void)z; + { + std::lock_guard g(SlotMutex()); + BufferSlot* s = FindSlot(p); + if (s == nullptr) return false; + s->device_current = true; + s->host_current = false; + } + return true; +} + + +// ITEM 5 (rope): driver-side warm hook. The decode-graph driver calls this +// for the step's (padded) positions BEFORE BeginCapture — the exact +// SizeSlot::Refresh slot in qwen3.cpp — so the persistent cos/sin tensors +// are populated outside capture and the captured rope cache-HITs on content. +// hq/hk select the expanded layouts to warm; base/args must match RopeNeox. +void WarmRopeCosSin(const int32_t* positions, int64_t tokens, int64_t hq, + int64_t hk, int64_t rot, double base) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + MeshDevice& device = SharedMeshDevice(); + std::vector cos_t, sin_t; + Tensor pos = Tensor::Contiguous(const_cast(positions), DType::kI32, + Device{DeviceType::kTENSTORRENT, 0}, {tokens}); + const RopeArgs no_scale{}; // plain rope only on the warm path + BuildCosSinFromPositions(pos, tokens, rot, base, no_scale, cos_t, sin_t); + // Byte-exact with what the captured rope reads: the per-step cos|sin CACHE + // stores f32-built values into a BF16 tensor (RopeCosSinCacheKernel's + // StoreElemF32 rounds), and the rope-side gather reads them back. Round the + // warm content through the same bf16 round-trip so the content-HIT + // comparison is exact. + for (auto& v : cos_t) v = BF16ToF32(F32ToBF16(v)); + for (auto& v : sin_t) v = BF16ToF32(F32ToBF16(v)); + auto warm_one = [&](int64_t heads) { + std::vector ce, se; + ExpandCosSinPerHead(cos_t.data(), sin_t.data(), tokens, heads, rot / 2, ce, se); + const uint32_t thu = static_cast(tokens * heads); + const uint32_t halfu = static_cast(rot / 2); + std::lock_guard g(RopeCSMutex()); + auto& c = RopeCSCache(); + const std::string k = RopeCSKey(thu, halfu); + auto it = c.find(k); + if (it == c.end() || it->second.cos_host != ce) { + // Populate OR refresh: the step's positions moved, so the persistent + // tensor's CONTENT must be updated (same address, new rows) — the + // SizeSlot::Refresh semantics. Legal here: the driver calls this + // outside capture. During capture the rope lookup then content-HITs. + RopeCSEntry e; + e.cos = UploadRows(ce.data(), thu, halfu, device); + e.sin = UploadRows(se.data(), thu, halfu, device); + e.cos_host = ce; + c[k] = std::move(e); + } + }; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmRopeCosSin tokens=%lld hq=%lld hk=%lld" + " rot=%lld first_pos=%d cos_first=%f\n", + (long long)tokens, (long long)hq, (long long)hk, + (long long)rot, (int)positions[0], + cos_t.empty() ? -1.0f : cos_t.front()); + warm_one(hq); + warm_one(hk); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) { + std::lock_guard g(RopeCSMutex()); + for (auto& [k, e] : RopeCSCache()) + std::fprintf(stderr, "[TT-TRACE] warm stored key=%s first=%f n=%zu\n", + k.c_str(), e.cos_host.empty() ? -1.0f : e.cos_host.front(), + e.cos_host.size()); + } +} + +void WarmPagedKvShadow(void* k_cache_data, void* v_cache_data, + int64_t num_blocks, int64_t block_size, + int64_t num_kv_heads, int64_t head_size, + int64_t used_blocks) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + if (num_blocks < 1 || block_size < 1 || used_blocks < 1) return; + MeshDevice& device = SharedMeshDevice(); + auto warm_one = [&](void* data) { + Tensor cache = Tensor::Contiguous( + data, DType::kBF16, Device{DeviceType::kTENSTORRENT, 0}, + {num_blocks, block_size, num_kv_heads, head_size}); + const uint32_t used = static_cast( + std::min(used_blocks, num_blocks)); + EnsurePagedKvTtnn(cache, device, used); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) { + std::lock_guard pg(PagedKvMutex()); + auto& sh = PagedKvShadows(); + std::fprintf(stderr, "[TT-TRACE] WarmPagedKvShadow ptr=%p nb=%lld used=%u shadows=%zu dev=%d\n", + data, (long long)num_blocks, used, sh.size(), + sh.count(reinterpret_cast(data)) ? + (int)sh[reinterpret_cast(data)].device.has_value() : -1); + } + }; + warm_one(k_cache_data); + warm_one(v_cache_data); +} + +void WarmRacIdx(const void* /*slot_mapping_owner*/, const int64_t* slots, + int64_t num_slots, int64_t block_size) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmRacIdx n=%lld bs=%lld slot0=%lld\n", + (long long)num_slots, (long long)block_size, (long long)slots[0]); + if (num_slots < 1) return; + MeshDevice& device = SharedMeshDevice(); + std::vector ptv; + std::vector idxv; + ptv.reserve(num_slots); idxv.reserve(num_slots); + for (int64_t t = 0; t < num_slots; ++t) { + const int64_t slot = slots[t]; + if (slot < 0) { ptv.push_back(0); idxv.push_back(0); continue; } + ptv.push_back(static_cast(slot / block_size)); + idxv.push_back(static_cast(slot % block_size)); + } + ttnn::Tensor pt = ttnn::Tensor::from_vector( + ptv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots), 1u}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + ttnn::Tensor idx = ttnn::Tensor::from_vector( + idxv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + std::lock_guard g(RacIdxMutex()); + RacIdxEntry e; + e.update_idxs = std::move(idx); + e.page_table = std::move(pt); + e.idx_host = idxv; + // Build the persistent sharded zero input [1,1,nkv_pad,d] (the shape + // paged_update_cache expects). Allocated OUTSIDE capture; the captured + // RAC copies the real k/v rows into its top slice then passes it. + // Build the persistent sharded zero input [1,1,nkv_pad,d] from the first + // available paged-KV shadow's geometry (same nkv/d as the cache). + { + std::lock_guard pg(PagedKvMutex()); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmRacIdx shadow loop: %zu shadows\n", + PagedKvShadows().size()); + for (auto& [ptr, shadow] : PagedKvShadows()) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] shadow ptr=%p nkv=%u d=%u dc=%d\n", + (void*)ptr, shadow.nkv, shadow.d, shadow.device_current); + if (shadow.nkv > 0 && shadow.d > 0) { + const uint32_t np = std::max(32u, ((shadow.nkv + 31u) / 32u) * 32u); + ttnn::Tensor z = ttnn::zeros( + ttnn::Shape({1u, 1u, np, shadow.d}), + ttnn::DataType::BFLOAT16, ttnn::Layout::TILE, std::ref(device)); + ttnn::Tensor reshaped = z.reshape(ttnn::Shape({1u, 1u, np, shadow.d})); + const auto grid = device.compute_with_storage_grid_size(); + const tt::tt_metal::CoreRangeSet core_set = + tt::tt_metal::num_cores_to_corerangeset(1u, grid, true); + tt::tt_metal::ShardSpec ss(core_set, {np, shadow.d}, + tt::tt_metal::ShardOrientation::ROW_MAJOR); + tt::tt_metal::MemoryConfig sm( + tt::tt_metal::TensorMemoryLayout::HEIGHT_SHARDED, + tt::tt_metal::BufferType::L1, ss); + e.sharded_zero = ttnn::to_memory_config(reshaped, sm); + // Also prime the zero cache for the padding-heads shape that + // build_input's capture branch uses (ZeroCacheGet on + // [1,1,pad_heads,d] TILE bf16). + const uint32_t ph = np - shadow.nkv; + if (ph > 0) { + ZeroCachePrime(ttnn::Shape({1u, 1u, ph, shadow.d}), + ttnn::DataType::BFLOAT16, ttnn::Layout::TILE, device); + } + break; + } + } + } + RacIdxCache()[std::make_pair(num_slots, block_size)] = std::move(e); + // paged_update_cache is warmed naturally: WarmPagedKvShadow (called by + // the driver BEFORE WarmRacIdx) primes the shadows, and the cold step's + // eager ForwardLayers runs TryReshapeAndCacheDeviceDecode (host_free is + // set, capturing is false) which calls paged_update_cache in the correct + // CQ context, compiling the program. +} + +void WarmPaMeta(const int32_t* block_table, int64_t num_reqs, int64_t max_blocks, + int64_t bt_row_stride, int64_t bt_col_stride, + const int32_t* seq_lens) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + if (num_reqs < 1) return; + MeshDevice& device = SharedMeshDevice(); + std::vector pt(static_cast(num_reqs * max_blocks)); + int32_t max_phys = -1; + for (int64_t r = 0; r < num_reqs; ++r) { + for (int64_t c = 0; c < max_blocks; ++c) { + const int32_t id = block_table[r * bt_row_stride + c * bt_col_stride]; + pt[static_cast(r * max_blocks + c)] = id; + if (id > max_phys) max_phys = id; + } + } + std::vector cpos(static_cast(num_reqs)); + for (int64_t r = 0; r < num_reqs; ++r) cpos[static_cast(r)] = seq_lens[r] - 1; + ttnn::Tensor dev_pt = ttnn::Tensor::from_vector( + pt, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs), + static_cast(max_blocks)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + ttnn::Tensor dev_cp = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + std::lock_guard g(PaMetaMutex()); + PaMetaEntry e; + e.page_table = std::move(dev_pt); + e.cur_pos = std::move(dev_cp); + e.pt_host = pt; + e.cp_host = cpos; + PaMetaCache()[std::make_pair(num_reqs, max_blocks)] = std::move(e); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmPaMeta n=%lld mb=%lld cp0=%d\n", + (long long)num_reqs, (long long)max_blocks, (int)cpos[0]); +} + +// Flush the pending RAC recorded during capture: write the recorded k/v +// device shadows into the KV cache shadows at the recorded slot. Called +// by the driver at the Refresh slot (outside capture). Uses the same +// paged_update_cache path as the eager RAC. +void FlushPendingRac() { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + std::vector entries; + int64_t slot = -1; + { + std::lock_guard g(PendingRacMutex()); + if (PendingRacVec().empty() || PendingRacSlotId() < 0) return; + entries = std::move(PendingRacVec()); + PendingRacVec().clear(); + slot = PendingRacSlotId(); + PendingRacSlotId() = -1; // consume + } + MeshDevice& device = SharedMeshDevice(); + // Get the RacIdxEntry for the idx tensors (warmed at Refresh). + std::optional rac; + { + std::lock_guard g(RacIdxMutex()); + if (!RacIdxCache().empty()) rac = RacIdxCache().begin()->second; + } + if (!rac.has_value()) return; + for (auto& ent : entries) { + void *kc = ent.k_cache, *vc = ent.v_cache; + ttnn::Tensor k_dev = ent.k_dev, v_dev = ent.v_dev; + // Get the paged-KV shadows for the recorded caches. + std::optional kc_dev, vc_dev; + uint32_t nkv = 0, d = 0; + { + std::lock_guard g(PagedKvMutex()); + auto* skc = &PagedKvShadows()[reinterpret_cast(kc)]; + auto* svc = &PagedKvShadows()[reinterpret_cast(vc)]; + if (skc->device.has_value() && svc->device.has_value()) { + kc_dev = skc->device; + vc_dev = svc->device; + nkv = skc->nkv; + d = skc->d; + } + } + if (!kc_dev.has_value() || !vc_dev.has_value() || nkv == 0) continue; + const uint32_t nkv_pad = std::max(32u, ((nkv + 31u) / 32u) * 32u); + auto build_input = [&](const ttnn::Tensor& src) -> ttnn::Tensor { + ttnn::Tensor reshaped = src.reshape(ttnn::Shape({1u, 1u, nkv, d})); + const uint32_t pad_heads = nkv_pad - nkv; + ttnn::Tensor padded; + if (pad_heads > 0) { + ttnn::Tensor zeros_tail = ttnn::zeros( + ttnn::Shape({1u, 1u, pad_heads, d}), + ttnn::DataType::BFLOAT16, ttnn::Layout::TILE, std::ref(device)); + padded = ttnn::concat(std::vector{reshaped, zeros_tail}, 2); + } else { + padded = reshaped; + } + const auto grid = device.compute_with_storage_grid_size(); + const tt::tt_metal::CoreRangeSet core_set = + tt::tt_metal::num_cores_to_corerangeset(1u, grid, true); + tt::tt_metal::ShardSpec ss(core_set, {nkv_pad, d}, + tt::tt_metal::ShardOrientation::ROW_MAJOR); + tt::tt_metal::MemoryConfig sm(tt::tt_metal::TensorMemoryLayout::HEIGHT_SHARDED, + tt::tt_metal::BufferType::L1, ss); + return ttnn::to_memory_config(padded, sm); + }; + ttnn::Tensor k_in = build_input(k_dev); + ttnn::Tensor v_in = build_input(v_dev); + ttnn::Tensor new_kc = ttnn::experimental::paged_update_cache( + *kc_dev, k_in, /*update_idxs=*/{}, rac->update_idxs, + /*share_cache=*/false, rac->page_table, 0, + std::nullopt, std::nullopt); + ttnn::Tensor new_vc = ttnn::experimental::paged_update_cache( + *vc_dev, v_in, /*update_idxs=*/{}, rac->update_idxs, + /*share_cache=*/false, rac->page_table, 0, + std::nullopt, std::nullopt); + { + std::lock_guard g(PagedKvMutex()); + PagedKvShadows()[reinterpret_cast(kc)].device = std::move(new_kc); + PagedKvShadows()[reinterpret_cast(kc)].device_current = true; + PagedKvShadows()[reinterpret_cast(vc)].device = std::move(new_vc); + PagedKvShadows()[reinterpret_cast(vc)].device_current = true; + } + } // per-entry loop + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] FlushPendingRac slot=%lld layers=%zu\n", + (long long)slot, entries.size()); +} } // namespace vt::tenstorrent diff --git a/tests/vt/test_tenstorrent_backend.cpp b/tests/vt/test_tenstorrent_backend.cpp index 8bb55668c..c34c63505 100644 --- a/tests/vt/test_tenstorrent_backend.cpp +++ b/tests/vt/test_tenstorrent_backend.cpp @@ -1520,3 +1520,65 @@ TEST_CASE("kTENSTORRENT kRmsNorm residual: device vs CPU f32 oracle across the r CHECK(max_abs < 0.05f); } } + +// BACKEND-TENSTORRENT-HOST-FREE-R1: guard the env-gated host-free helpers' +// DEFAULT-PATH INERTNESS. The helpers (CopyDeviceDeviceIfCapture / +// MemsetDeviceIfCapture, vt/tenstorrent/tenstorrent_device.h) must DECLINE +// unless VT_TT_HOST_FREE_DECODE is set (or capture is active). Without this +// case that property is enforced by code review alone: a removed gate flips +// ordinary eager Copy/Memset to device variants silently (review mutation M1) +// and a capture flag stuck true after a failed EndCapture does the same (M4). +// Both buffers below carry CURRENT device shadows with equal byte sizes, so +// the flag gate is the ONLY thing that can make the helpers decline. +#include "../../src/vt/tenstorrent/tenstorrent_device.h" + +TEST_CASE("kTENSTORRENT host-free helpers decline by default (inertness guard)") { + if (!TenstorrentPresent()) { + MESSAGE("SKIPPED: no Tenstorrent device on this box"); + return; + } + ::unsetenv("VT_TT_HOST_FREE_DECODE"); // the guard is about the UNSET case + Backend& backend = *vt::TryGetBackend(DeviceType::kTENSTORRENT); + + // Two same-shaped outputs, each given a current device shadow by a device + // Matmul (CommitDevice2D leaves device_current=true, host_current=false). + constexpr int64_t M = 8, K = 32, N = 8; + auto shadowed = [&](std::vector& host) { + std::vector a(M * K, 0.5f), b(K * N, 0.25f); + host.assign(static_cast(M * N), -1.0f); + void* ma = backend.Alloc(a.size() * sizeof(float)); + void* mb = backend.Alloc(b.size() * sizeof(float)); + void* mo = backend.Alloc(host.size() * sizeof(float)); + Queue q = backend.CreateQueue(); + backend.Copy(q, ma, a.data(), a.size() * sizeof(float)); + backend.Copy(q, mb, b.data(), b.size() * sizeof(float)); + Tensor ta = Tensor::Contiguous(ma, vt::DType::kF32, Device{DeviceType::kTENSTORRENT, 0}, {M, K}); + Tensor tb = Tensor::Contiguous(mb, vt::DType::kF32, Device{DeviceType::kTENSTORRENT, 0}, {K, N}); + Tensor to = Tensor::Contiguous(mo, vt::DType::kF32, Device{DeviceType::kTENSTORRENT, 0}, {M, N}); + reinterpret_cast(vt::GetOp(vt::OpId::kMatmul, DeviceType::kTENSTORRENT))(q, to, ta, tb); + return mo; // caller keeps the allocation; shadow lives in the slot map + }; + std::vector h1, h2; + void* m1 = shadowed(h1); + void* m2 = shadowed(h2); + + // The gate: same bytes, both shadows current -> only the env/capture gate + // can decline. These CHECKs go RED if the gate is removed (M1) or if the + // capture flag is stuck true (M4). + CHECK_FALSE(vt::tenstorrent::CopyDeviceDeviceIfCapture(m2, m1)); + CHECK_FALSE(vt::tenstorrent::MemsetDeviceIfCapture(m2, 0)); + // value!=0 always declines (host memset is the only path for it). + CHECK_FALSE(vt::tenstorrent::MemsetDeviceIfCapture(m2, 1)); + + // And the default host path still works: Copy m1 -> m2 yields identical + // host bytes once materialized. + Queue q = backend.CreateQueue(); + std::vector got(h1.size(), -7.0f); + backend.Copy(q, m2, m1, h1.size() * sizeof(float)); + backend.Copy(q, got.data(), m2, got.size() * sizeof(float)); + // 0.5f * 0.25f summed over K=32 == 4.0f per element (bf16 device acc). + CHECK(got == std::vector(static_cast(M * N), 4.0f)); + + backend.Free(m1); + backend.Free(m2); +}