diff --git a/.agents/backend-matrix.md b/.agents/backend-matrix.md index 74a5e4ee9..b6307c4e1 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/issue-index.md b/.agents/issue-index.md index d47181eec..4d24d5746 100644 --- a/.agents/issue-index.md +++ b/.agents/issue-index.md @@ -343,3 +343,4 @@ rather than merged. `scripts/check-agent-record.py` gates both. | [#1144](https://github.com/mudler/vllm.cpp/issues/1144) | — | `Res2sTwoStageRecipe` is this tree's port of `TI2VidTwoStagesHQPipeline` (`ti2vid_two_stages_hq.py:59` @ `fd4ded7f`), the one in-scope pipeline that runs the SAME adapter on BOTH stages at DIFFERENT strengths: it builds `distilled_lora_stage_1` and `distilled_lora_stage_2` from one path (`:92-101`) and hands one to each `DiffusionStage.from_checkpoint` (`:154`, `:165`), CLI-defaulted 0.25 and 0.5 (`utils/args.py:1174-1184`). Here neither phase sets `Ltx2PhaseRecipe::loras`, so both take the `kAllAdapters` default, and the engine carries ONE strength for the whole load (`lora_strength` absent is 1.0, `include/vllm/multimodal/ltx2_video.h:214-218`) — so both stages run at 1.0. Nothing refuses and nothing changes shape; only a render against upstream on the same checkpoint, take and seed sees it. A NEW FIELD ON `Ltx2PhaseRecipe` DOES NOT CLOSE IT: `Ltx2RebindDitLoras` early-returns on `currently_fused == fuse` where `currently_fused` is `checkpoint.lora_fused_tensors > 0`, a BOOLEAN, so it detects "already fused" and never "already fused AT THIS STRENGTH" — and HQ has both stages fused, so the no-op its own header advertises would swallow the transition and stage 2 would render at stage 1's strength. Closing it needs `bool fuse` to become a type carrying a strength AND `Ltx2DitCheckpoint` to record WHICH adapter state is applied, plus a `--distilled-lora-strength-stage-1`/`-stage-2` request pair, since one `lora_strength` extra cannot spell two values. The re-materialize-and-write-back mechanism #1118 landed is untouched by that change. Filed because [#921](https://github.com/mudler/vllm.cpp/issues/921), which owned the per-phase strength in [`ltx25-phase-lora.md`](specs/ltx25-phase-lora.md), was CLOSED as completed on 2026-08-17 by `LTX25-RES2S-LOOP` (`4d7748646`, PR [#1125](https://github.com/mudler/vllm.cpp/pull/1125)); that row named the distilled LoRA per stage as out of scope and correct to leave, but did not list it under its own `## Owed`, so the debt outlived its issue with no open owner. Found during the review repair of PR [#1140](https://github.com/mudler/vllm.cpp/pull/1140). Listed under `## Owed` in [`ltx25-phase-lora.md`](specs/ltx25-phase-lora.md) | bug | | [#584](https://github.com/mudler/vllm.cpp/issues/584) | `ENG-RELEASE-WINDOWS` | `test_openai_api_server.exe` fast-fails with `-1073740791` / `0xC0000409` on BOTH Windows lanes, and the whole doctest output is the version banner — no `Status:` line, no `assertions:` line. NARROWED here while landing [#503](https://github.com/mudler/vllm.cpp/issues/503), not fixed. `0xC0000409` is NOT evidence of a stack buffer overrun: it is the status `__fastfail` raises for every fail-fast code, so `abort()` — and therefore `std::terminate()` — and the CRT invalid-parameter handler both surface as it, and `__fastfail` bypasses SEH, which is why doctest's Windows handler cannot report it. LOCALISED by the log: `LogHttpIngress` (`src/vllm/entrypoints/openai/api_server.cpp:223`) fires for `/v1/chat/completions` and no other route, over `std::cerr` (`request_logger.cpp:26`, unit-buffered, so absence is evidence and not buffering); `tests/vllm/entrypoints/openai/test_api_server.cpp` reaches that route at exactly eight sites in file order (`:596 :647 :662 :703 :748 :781 :902 :1292`), and the eighth and last logged request (`body_bytes=92 stream=0 max_tokens=4 prompt_chars=5`) is byte-for-byte the 92-byte body posted at `:1291-1294` inside `TEST_CASE("api_server: socket smoke — real HTTP requests over an ephemeral port")`. The server ANSWERED it (`Finished request chatcmpl-0 completion_tokens=4`), then produced nothing for 0.78 s and fast-failed — which places the fault between `:1294` and the `REQUIRE` at `:1325`, and excludes the 1.0 s poll loop at `:1323-1324`. WHY IT IS UNDIAGNOSABLE, and this half is provable by inspection rather than inferred: fifteen cases hold a joinable `std::thread` across throwing assertions (`:1243/1303` … `:3159/3222`, plus two threads in the `_WIN32`-only teardown case at `:2616`), so ANY throw in between destroys a joinable thread → `std::terminate()` → `abort()` → `__fastfail`, turning a named assertion failure into an opaque `0xC0000409` with no reporter output. NOT fixed in flow: it is a C++ edit at fifteen sites in a file that session could not build (host at 94 % disk; a mutation that fails to compile reads as a passing test), and landing it unbuilt would risk the green Linux lanes to repair a Windows instrument. Deterministic across four unrelated trees (`76f2a6d84`, `c83b96934`, `04b58bf6a`, the #1069 merge `6314602bb`) and identical on `windows-msvc-vulkan`, so both lanes are blocked by ONE defect. Evidence: jobs `95490912332`, `95479509343`, `95327871315`, `95479509369`. Diagnosis and next step under `## Owed` in [`windows-baseline-coverage.md`](specs/windows-baseline-coverage.md) | bug | | [#1148](https://github.com/mudler/vllm.cpp/issues/1148) | `LTX25-BF16-DIT` | `PlanDit` refused every DiT checkpoint carrying neither `U8` nor `F8_E4M3` (`src/vllm/model_executor/models/ltx2_loader.cpp:414-418` @ `c83b96934`), so `ltx-2.5-22b-dev-transformer-bf16.safetensors` — the FULL model, 42,018,190,584 bytes, 4349 tensors, 4059 BF16 / 290 F32, ZERO `_scale` names, measured from its own header on 2026-08-17 — could not be read at all. Upstream's pipeline table (`packages/ltx-pipelines/CLAUDE.md:17-30` @ `fd4ded7f`) names that model for `TI2VidOneStagePipeline`, `T2AOneStagePipeline`, `TI2VidTwoStagesPipeline`, `TI2VidTwoStagesHQPipeline`, `A2VidPipelineTwoStage` and `KeyframeInterpolationPipeline`, so `one_stage`, `t2a_one_stage`, `res2s_two_stage` and `a2vid_two_stage` were all landed and all runnable only against a DISTILLED checkpoint, which is a different sampling regime that renders plausibly. The refusal's own advice was unreachable: it said "use the L2 path" and `Ltx2LoadDitFromSafetensors` IS the L2 path, calling `PlanDit` on its first line, as do `Ltx2ParseDitParamsFromCheckpoint`, `Ltx2StreamDitToDevice` and `Ltx2RebindDitLoras`. Upstream has no third quant state — `_DTYPE_CASTABLE` (`single_gpu_model_builder.py:51-57`) lists float32/float64/float16/bfloat16 and calls uint8-NVFP4 and float8 "quantized payloads", so unquantized is the BASELINE the two arms are exceptions to — and `MaterializeDitTensor` already carried a `BF16` branch, so the wall was the one `if` and nothing behind it. Fixed by `Ltx2DitQuant::kNone`, with the refusal that survives naming the dtypes the file holds. Spec [`ltx25-bf16-dit.md`](specs/ltx25-bf16-dit.md) | bug | +| [#1105](https://github.com/mudler/vllm.cpp/issues/1105) | `BACKEND-TENSTORRENT-HOST-FREE-FORWARD` | Tenstorrent host-free decode graph: capture, replay, and on-device `cur_pos` advance. Owns the trace-runner SPIKE that measured the `to_vector` capture abort and the ACTIVE host-free forward row. Spec [`tenstorrent-host-free-forward.md`](specs/tenstorrent-host-free-forward.md) | feature | diff --git a/.agents/specs/tenstorrent-host-free-forward.md b/.agents/specs/tenstorrent-host-free-forward.md new file mode 100644 index 000000000..1cbc0f59e --- /dev/null +++ b/.agents/specs/tenstorrent-host-free-forward.md @@ -0,0 +1,289 @@ +# 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`). Issue: +[#1105](https://github.com/mudler/vllm.cpp/issues/1105). + +## 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). + +## Owed + +- **`DecodePosCache` is keyed on bare `num_reqs`, with no engine, queue or device + identity, and is never cleared.** Two engine instances in one process at the same + padded batch size therefore share one `cur_pos` device tensor: the second engine's + first `WarmDecodePos` finds the first engine's entry, returns early, and both + `WarmRacIdx` / `WarmPaMeta` aliases bind to a buffer another engine is advancing. + That is silently wrong rather than a refusal, and it is a candidate explanation for + `test_qwen3_paged_engine` still timing out under the flag. Owned by + [#1105](https://github.com/mudler/vllm.cpp/issues/1105). +- **`VT_TT_RECAPTURE_EVERY` lags `cur_pos` by one per recapture cycle.** + `GraphCapturesCounter` is only ever `fetch_add`ed (`tenstorrent_ops.cpp:3225`, called + once from `:3297`); nothing resets it, and `DestroyGraph` does not. So the + recapture-triggered eager step re-seeds nothing and skips every `copy_to_device`, + and the following capture step captures a `cur_pos` one behind. The eager PA + consistency check at `:2410-2412` cannot see it: it compares `e.cp_host[0]` against + the same `seq_lens` that rewrote `e.cp_host` at `:3982`, so it validates the host + mirror against itself. Owned by [#1105](https://github.com/mudler/vllm.cpp/issues/1105). + +The seven constraints above remain. A new batch size after the first +capture is now refused (`VT_CHECK` in `WarmDecodePos` / `WarmPaMeta` / +`WarmRacIdx`) rather than freezing `cur_pos`. The real fix is a +per-cache-entry seed / generation field and aliasing on every warm call, +not a process-global `GraphCapturesCounter`. Tracked on +[#1105](https://github.com/mudler/vllm.cpp/issues/1105). + +The operator must still rerun the 80-token no-hang gate and +`test_qwen3_paged_engine` on a Blackhole P150. An implementer run is an +input, not a gate result. + +## Now + +`ACTIVE`. R1-R3b and the R2 on-device `cur_pos` / `update_idxs` advance are +implemented on this branch, env-gated by `VT_TT_HOST_FREE_DECODE`. A P150 +run of Qwen3-0.6B "Hello" at 80 tokens completed 79 replays with no hang +and 22/22 argmax vs the per-step-copy baseline. Next: operator rerun of +that 80-token gate and `test_qwen3_paged_engine` on card. A new batch +size after the first capture now throws instead of emitting wrong tokens. diff --git a/.agents/specs/tenstorrent-host-free-r1.md b/.agents/specs/tenstorrent-host-free-r1.md new file mode 100644 index 000000000..f071c93cd --- /dev/null +++ b/.agents/specs/tenstorrent-host-free-r1.md @@ -0,0 +1,1158 @@ +# 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`. + +## Now + +`ACTIVE` on the parent row `BACKEND-TENSTORRENT-HOST-FREE-FORWARD`. R1-R3b +and the R2 on-device `cur_pos` advance are on this branch. Next: fresh review +of the host-free decode graph, then the operator reruns the Qwen3-0.6B +80-token no-hang gate and the TT golden on a Blackhole P150. + +## 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. + +## Session checkpoint (2026-08-16) — host-free decode graph fidelity + stall characterization + +### Complete state + +**HOST-FREE DECODE GRAPH LANDS.** Captured replay is token-exact vs +host-free eager (22/22 argmax-identical, per-step `VT_TT_DUMP_KV` diff, +monitor-57 vs monitor-58). Marginal decode step: captured ~30 ms (~33 +tok/s) vs default eager ~175 ms (~5.7 tok/s) ⇒ **5.8× speedup** at the +step level (short-prompt startup ~21.4 s still dominates wall time). + +### What works (all env-gated, inert by default) + +- Frozen-replay fix: `InvalidateHostCachesAfterTrace()` after every + `execute_trace` (in `TraceReplay`/`TraceReplayGraph`, both + `blocking=false` — the qwen3 driver's blocking logits readback is the + sync point, mirroring upstream executor.py's traced-decode pattern). +- KV shadow pool-sized allocation: `EnsurePagedKvTtnn` rewritten — used + prefix `from_vector` + `ttnn::zeros` tail + `ttnn::concat`; the shadow + buffer is allocated ONCE at the full pool size and never moves (the old + per-block-boundary realloc freed a buffer a live trace still + referenced → all-zero logits at the first block-table growth). +- Embedding moved INSIDE the captured trace: `WarmDecodeIds` (persistent + device UINT32 ids tensor, content refreshed in place, zero eager + allocs per replay) + `EmbedDeviceIdsInto` (capture-safe embedding over + the persistent ids tensor writing the persistent hidden shadow). A + replay step now performs ZERO eager device allocations. +- Re-capture workaround `VT_TT_RECAPTURE_EVERY=N`: destroy + re-capture + the graph every N replays (see stall, below). + +### The two walls (characterized + cause ISOLATED 2026-08-16) + +1. **~38-replay completion hang.** A single trace id deterministically + stops completing (futex wait in the post-replay readback) after ~38–50 + consecutive replays. `execute_trace` enqueues fine (0.08 ms); the + device-side completion never arrives. +2. **~2 healthy capture generations.** `VT_TT_RECAPTURE_EVERY=N`: + N=32 → gen1=32 ✓, gen2=32 ✓, gen3 died @11 (~75 replays total). + +**Upstream machinery is NOT the cause** (monitors 61-66). Ported upstream +`test_mesh_trace.cpp` Sanity + 4 scratch variants (120 replays of one trace, +Finish-per-replay, 10 recapture gens, 50-program big trace) — ALL PASS on +our card. The raw mesh-trace path is sound. + +**Cause ISOLATED by per-class bisection** (monitors 68-72): + +| skip flag(s) ON | copies active | result | replays | +|---|---|---|---| +| (none, recapture=32) | rope+idx+PAmeta+ids | HANG | 75 | +| `VT_TT_NO_ROPE_REFRESH` | idx+PAmeta+ids | HANG | 39 | +| `VT_TT_NO_IDX_WARM` | rope+ids | **PASS** | 79 | +| `VT_TT_NO_IDS_WARM` | rope+idx+PAmeta | HANG | 50 | +| all three | (none) | **PASS** | 79 | + +The toxic copy class is `WarmRacIdx` + `WarmPaMeta` (the +`VT_TT_NO_IDX_WARM` flag), both calling `ttnn::copy_to_device` into INT32 +`page_table`/`update_idxs`/`cur_pos` buffers the captured trace READS. +Rope (bf16) and decode-ids (UINT32) are innocent. + +`ttnn::copy_to_device` (ttnn/core/tensor/tensor_ops.cpp:182) → +`enqueue_write_tensor`: a DIRECT host→device write into the existing +device buffer on the trace's CQ, NO allocation. So the leak is NOT an +allocation (the `mark_allocations_unsafe` warning correctly never fired) — +it's a write to a trace-input buffer on the trace's CQ between replays. + +**Root mechanism (code-read 2026-08-16):** `copy_to_device` → mesh CQ +`write_to_core` (fd_mesh_command_queue.cpp:579) → +`device_dispatch::write_to_core`, which threads +`expected_num_workers_completed` (impl/buffers/dispatch.cpp:486+, +737/985/1304/1350) — the SAME counter the trace replay path +(`update_worker_state_post_trace_execution`, tt_metal/impl/trace/dispatch.cpp) +assigns and `mark_completely_full()`s per replay. After ~38–50 interleaved +writes, the accounting desyncs and the next replay's device-side completion +never arrives → the post-replay readback futex-hangs. + +The toxic class is specifically idx/PAmeta: with idx/PAmeta OFF (rope+ids +ON) the run PASSES 79 replays; with idx/PAmeta ON it hangs regardless of +the others. The wall count (39/50/75) is run-to-run variance, not +evidence that rope/ids contribute. Why idx/PAmeta specifically (not rope +bf16, not ids UINT32) is not yet distinguished — candidates: they're read +by `sdpa_decode`/`paged_update_cache` dispatch (prefetch dependency), or +the INT32 ROW_MAJOR 2D page_table write is the specific dispatch shape +that leaks. + +Memory notes: `tt-trace-replay-wall-not-upstream-machinery.md`, +`tt-replay-wall-toxic-copy-class-idx-pameta.md`, +`tt-replay-wall-root-mechanism-write-to-core.md`. + +### Performance (Qwen3-0.6B, "Hello", bf16, 1 req) + +| config | step | marginal tok/s | +|--------|------|----------------| +| default eager | ~175 ms | ~5.7 | +| captured replay | ~30 ms | ~33 | +| speedup | — | 5.8× | + +Fixed startup ≈21.4 s dominates short runs. Re-capture cycle <1 s. + +### Mode difference (NOT a capture bug) + +Host-free mode (eager or captured) vs default mode diverge at generation +step 18 (host-free emits " first." where default emits "."), re-sync +after. Both coherent; near-tie flip caused by host-free ops changes +(pool-sized KV shadow, padded-view paths). The golden-gate near-tie +arrays are the arbiter — not yet run. + +### Files modified this session (all uncommitted) + +- `src/vt/tenstorrent/tenstorrent_ops.cpp`: + `InvalidateHostCachesAfterTrace()` + calls; `EnsurePagedKvTtnn` + rewrite; `DecodeIdsCache`/`WarmDecodeIds`/`EmbedDeviceIdsInto`; + `GraphCapturesCounter`/`Done`/`NoteGraphCaptured`/`ReplayRegimeBisectSkip`; + stall-bisection skip flags (`VT_TT_NO_ROPE_REFRESH`/`VT_TT_NO_IDX_WARM`/ + `VT_TT_NO_IDS_WARM`, replay-regime only); `TraceDestroyGraph` + `release_trace` diagnostics; `[TT-STEP]` brackets in `TraceReplayGraph`. +- `src/vt/tenstorrent/tenstorrent_device.h`: declarations for + `WarmDecodeIds`, `EmbedDeviceIdsInto`. +- `src/vllm/model_executor/models/qwen3.cpp`: `VT_TT_RECAPTURE_EVERY` + driver logic; replay-branch timing (`replay_ns`/`replay_steps` on Impl); + `~Impl()` stats extension. + +### What's left + +1. **Fix the toxic copy path** (idx/PAmeta `copy_to_device` into trace-input + buffers). Fix directions (ranked): (a) on-device state advance inside the + trace (upstream executor.py pattern, biggest change); (b) capture-safe + write primitive (tiny captured program); (c) raw `enqueue_write_buffer` + bypass of ttnn's tensor write path; (d) separate CQ. Correctness gate + (needs >75 replays) is blocked until a fix or longer-lived workaround. +2. Fresh review of the complete host-free change; commit everything + (nothing committed yet — all changes uncommitted). +3. Upstream issue to tt-metal: the ~38-replay wall under our ttnn-op + interleaving + the ~2-generation re-capture wall + the earlier + padded-shape view trap + the TILE 4D view mapping rule. diff --git a/.agents/specs/tenstorrent-host-free-r2.md b/.agents/specs/tenstorrent-host-free-r2.md new file mode 100644 index 000000000..142e970fb --- /dev/null +++ b/.agents/specs/tenstorrent-host-free-r2.md @@ -0,0 +1,198 @@ +# Spike: on-device decode state advancement (host-free R2) + +**Status:** SPIKE (pre-implementation spec) — POL-SPIKE-FIRST. +**Branch:** `row/BACKEND-TENSTORRENT-HOST-FREE-R2` (off `row/BACKEND-TENSTORRENT-HOST-FREE-R1`). +**Owning spec:** `.agents/specs/tenstorrent-host-free-r1.md` (this extends the +host-free decode work; R1 is the captured-replay baseline). + +## Problem + +R1's captured decode graph is token-exact and 5.8× faster than eager, but +hangs deterministically after ~38–50 consecutive replays of one trace id +(monitor-68). Bisection (monitors 68–72) isolated the cause to the per-replay +**eager host→device copies** into trace-input INT32 buffers: + +- `WarmRacIdx`: `copy_to_device` into `update_idxs` + `page_table` (INT32) +- `WarmPaMeta`: `copy_to_device` into `page_table` + `cur_pos` (INT32) + +With these OFF (`VT_TT_NO_IDX_WARM`), 79 replays complete cleanly; with them +ON, hang at ~38–50. Root mechanism (code-read): `copy_to_device` → mesh CQ +`write_to_core` perturbs `expected_num_workers_completed`, the same dispatch +counter the trace replay path assigns/`mark_completely_full`s per replay. + +**A minimal standalone repro was NOT found** (monitors 73–90): the toxic +write pattern + the exact RAC+sdpa ops, at 1× and 28×, with `blocking=False` +and sharded L1 inputs, all PASS 120 replays. The trigger is specific to the +full diverse-layer driver trace. So the fix must eliminate the interleaved +writes entirely rather than work around a cleanly-reproducible upstream bug. + +## Goal + +Eliminate the per-replay eager host→device copies of decode state +(`cur_pos`, `update_idxs`, rope cos/sin, decode ids) by advancing that state +**on-device inside the captured trace**, mirroring the upstream +`executor.py` traced-decode pattern (`models/common/models/executor.py:2186`). +After this change, a steady-state replay step should perform **0 eager +host→device copies** (except `page_table`, only when a block boundary is +crossed — rare, not every step). + +## Upstream pattern (the reference) + +`executor.py:2186` comment: *"the captured trace advanced current_pos/ +rot_mat_idxs (in-place plus_one) and wrote the sampled token back into the +persistent token buffer on the previous replay, so tokens/positions need NO +host staging. Refresh only the page table, and only when it actually changed +(block boundaries crossed)."` + +Concrete (`executor.py:2310-2351`): +1. Persistent device tensors: `tt_tokens`, `tt_current_pos`, `tt_rot_mat_idxs`, + `tt_page_table` (allocated once before capture). +2. Captured body: `embed(tokens)` → `decode_forward(..., current_pos, rot_mats, + page_table)` → (on-device sampling writes token back to `tt_tokens`) → + `ttnn.plus_one(tt_current_pos)` → `ttnn.plus_one(tt_rot_mat_idxs)`. +3. Steady-state replay: `copy_host_to_device_tensor(page_table, tt_page_table)` + only if `page_table` changed since last step; `execute_trace(blocking=False)`; + read back logits/tokens. + +`ttnn::operations::experimental::plus_one` exists in our checkout +(`ttnn/cpp/ttnn/operations/experimental/plusone/plusone.hpp:14`) — in-place, +traceable, with `skip_negative_entries`. + +## Our driver's gap + +Our driver samples on host (downloads logits → CPU argmax → next token), so +the `tt_out_tok` in-place-writeback (which needs on-device sampling) does not +directly apply. But the **position/idx advancement** — the actual toxic class +— is independent of sampling and IS movable on-device: + +- `cur_pos = seq_lens - 1`, advances +1 per step → `plus_one` on a persistent + `cur_pos` tensor captured in the trace. +- `update_idxs = seq_lens - 1` (same value as cur_pos for decode T=1) → derive + from `cur_pos` on-device (alias or `plus_one` a parallel persistent tensor). +- rope cos/sin indexed by `positions` (= cur_pos) → upstream uses + `ttnn.plus_one(tt_rot_mat_idxs)`; our `WarmRopeCosSin` recomputes from + `positions` each step. Movable: capture a `plus_one` on the rope index, or + capture the rope lookup itself. +- `decode ids` (the sampled token for embedding) — **proven innocent by + bisection** (`VT_TT_NO_IDS_WARM` alone still hangs; ids OFF + idx ON hangs). + Still, eliminating it closes the loop: the host-sampled token must reach the + embedding input. Since we sample on host, this one host→device copy per step + is unavoidable UNLESS we move sampling on-device. **Defer** — ids is not the + toxic class, keep the host copy for now (1 copy/step is fine, matching the + upstream "1 copy when page_table changes" budget). +- `page_table` → refresh from host only when a block boundary is crossed + (block-table grows). Rare; not every step. + +## Design + +### Phase 1 (the toxic class): on-device `cur_pos` + `update_idxs` + +1. **Persistent `cur_pos` tensor** (INT32, `[B]`, device), allocated once before + capture. Seeded from host at capture time (first decode step's `seq_lens-1`). +2. **Capture `ttnn::plus_one(cur_pos, skip_negative_entries=true)` at the END of + the trace body** (after all reads of `cur_pos` in sdpa_decode/RAC), so the + NEXT replay sees the incremented value. Warm the `plus_one` program before + capture (program-cache). +3. **`update_idxs`**: for decode T=1, `update_idxs == cur_pos`. Two options: + (a) alias — pass the same persistent tensor as `update_idxs_tensor` to + `paged_update_cache` (if the op accepts it without copy); (b) a second + persistent tensor `+1`'d in parallel. Prefer (a); fall back to (b). +4. **`WarmPaMeta`/`WarmRacIdx` become no-ops in steady state** — they only fire + on re-seed (first step after prefill, or after a re-capture) to set the + initial `cur_pos` from host. Gate them on `!replay_regime` (i.e. only on the + capture/warm step), not every step. + +### Phase 2 (page_table): refresh only on change + +5. Track `prev_page_table` (host-side). On each step, `if (page_table != + prev_page_table) copy_host_to_device(page_table, tt_page_table)`. Block + boundaries are rare (every `block_size` tokens), so this is ~0 copies/step + in steady state. + +### Phase 3 (rope, deferred) + +6. `WarmRopeCosSin` is **not** the toxic class (bisection: rope OFF still hangs + at 39; idx OFF passes). Defer moving rope on-device to a later iteration + unless Phase 1+2 don't clear the wall (they should — idx/PAmeta is the sole + trigger). + +### Phase 4 (decode ids, deferred) + +7. `WarmDecodeIds` is innocent (ids OFF + idx ON still hangs). Keep the host + copy for the sampled token (1 copy/step). Moving sampling on-device is a + larger change and not needed to clear the wall. + +## Done-when + +- A steady-state replay step performs 0 eager host→device copies of + `cur_pos`/`update_idxs`/rope/ids (only `page_table` on block-boundary steps). +- `VT_TT_HOST_FREE_DECODE=1` decode of Qwen3-0.6B "Hello" `--max-tokens 80` + completes 79 replays with **no hang** (currently hangs at ~38–50 with idx + warm ON, passes only with `VT_TT_NO_IDX_WARM`). +- Token-exact fidelity preserved: 22/22 argmax vs host-free eager (re-run the + m57-vs-m58 diff protocol). +- No new eager device allocations during replay (the existing + `InvalidateHostCachesAfterTrace` + pool-sized shadow invariants hold). + +## Evidence / gates + +- Focused gate: the 80-token no-hang run above + the fidelity diff. +- Full gate: `scripts/agent-preflight.sh` + the TT golden correctness gate + (unblocks once >75 replays work — currently blocked by the wall). +- Fresh review + operator rerun per POL-REVIEW-FRESH / POL-OPERATOR-VERIFY. + +## Risks + +- `plus_one` may not be capture-safe in the way we need (in-place write to a + buffer the trace also reads earlier). Upstream's comment ("These writes are + AFTER every read of the buffers in this body, so there is no intra-trace + hazard") suggests it's fine, but verify. +- `paged_update_cache` may not accept an aliased `update_idxs_tensor` (it may + copy internally). If so, Phase 1(b) needs a second persistent tensor + a + second `plus_one`. +- Moving `cur_pos` on-device changes the R1 captured-graph semantics; the + fidelity diff must re-confirm token-exactness. +- This is a larger change than R1's incremental fixes; spec-first, red-test- + before-fix, focused-then-full gate per policy. + +## Stop conditions + +- If Phase 1 (on-device cur_pos/update_idxs) does not clear the wall (i.e. the + hang persists with idx/PAmeta copies eliminated), the trigger is NOT the + copies but something else in the full driver trace — re-bisect with the new + minimal copy budget, and reconsider the upstream issue. +- If `plus_one` cannot be captured safely, fall back to a host-side + `copy_to_device` of `cur_pos` only (1 copy/step) and re-test — if THAT + clears the wall, the trigger was the multi-copy churn, not the write per se. + +## Now + +`ACTIVE` as the implementation wave of `BACKEND-TENSTORRENT-HOST-FREE-FORWARD`. +On-device `cur_pos` plus_one is captured in the decode graph. A P150 run of +Qwen3-0.6B "Hello" at 80 tokens completed 79 replays with no hang. Next: +fresh review, then the operator reruns that gate and the TT golden on card. + +## First resume command + +```sh +cd /home/lu_zero/Sources/vllmcpp-tenstorrent +git status # confirm clean base on row/BACKEND-TENSTORRENT-HOST-FREE-R1 +# Read this spec, then implement Phase 1: persistent cur_pos + plus_one in trace. +``` + +## Probe results (2026-08-16, monitor-92) + +**`ttnn.plus_one` is capture-safe and advances on-device on replay.** +Probe (`scripts/probe_plus_one_capture.py`): persistent INT32 `cur_pos=[5]`, +warm `plus_one` → `[6]`, capture a trace whose body is just `plus_one(cur_pos)`, +replay 5× → `[7,8,9,10,11]`. Clean increment, no capture fatal, no hang. + +- Python binding: `ttnn.plus_one` (top-level, NOT `ttnn.experimental.plus_one`). +- C++ API: `ttnn::operations::experimental::plus_one(tensor, sub_core_grids, + skip_negative_entries)` in + `ttnn/cpp/ttnn/operations/experimental/plusone/plusone.hpp:14`. +- In-place, traceable, with `skip_negative_entries` (matches upstream + `plus_one(current_pos, skip_negative_entries=True)`). + +⇒ Phase 1 risk #1 (capture-safety) is cleared. Phase 1 risk #2 (aliased +`update_idxs_tensor` in `paged_update_cache`) still to probe during impl. diff --git a/.agents/specs/tenstorrent-trace-runner.md b/.agents/specs/tenstorrent-trace-runner.md new file mode 100644 index 000000000..b39abc7de --- /dev/null +++ b/.agents/specs/tenstorrent-trace-runner.md @@ -0,0 +1,271 @@ +# 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. + +## Now + +`SPIKE`. NO-GO for pure T=1 decode capture: the T=1 forward does `to_vector` +readbacks that ttnn trace prohibits. Decode capture moved to +`BACKEND-TENSTORRENT-HOST-FREE-FORWARD` (#1105). Next for this row: audit +`to_vector` in the prefill / multi-token region before claiming capture +there. diff --git a/.agents/specs/tt-metal-trace-replay-write-desync.md b/.agents/specs/tt-metal-trace-replay-write-desync.md new file mode 100644 index 000000000..69e08fdd5 --- /dev/null +++ b/.agents/specs/tt-metal-trace-replay-write-desync.md @@ -0,0 +1,180 @@ +# Interleaved host→device writes between trace replays desync `expected_num_workers_completed` → completion hang after ~38–50 replays + +## Summary + +On a single Blackhole P150 (1×1 mesh), replaying a captured mesh trace that +captures ttnn ops (embedding + sdpa_decode + paged_update_cache, ~50 MB trace) +hangs deterministically after ~38–50 consecutive replays of one trace id. The +hang is a futex wait in the post-replay blocking readback: `execute_trace` +returns (the trace is enqueued) but the device-side completion never arrives. + +**Root cause (isolated by bisection + code read):** issuing +`ttnn::copy_to_device` (→ `enqueue_write_tensor` → mesh CQ `write_to_core`) +into a buffer the captured trace *reads*, between replays while the trace is +live, perturbs `expected_num_workers_completed` — the same dispatch counter the +trace replay path (`update_worker_state_post_trace_execution`) assigns and +`mark_completely_full()`s per replay. After enough interleaved writes the +accounting desyncs and the next replay's `DISPATCH_WAIT(count)` never +completes. This is the between-replays analogue of the within-trace bug fixed +in #7978 / discussed in #7793. + +## Environment + +- tt-metal: `a3d330289752192754277638fe5c09eb2fb49763` (2026-08-09) +- Device: single Blackhole P150 (1×1 mesh, `trace_region_size=50 MB`, + `l1_small_size=DEFAULT_L1_SMALL_SIZE`, 1 CQ) +- Firmware bundle: 19.7.1 +- The captured trace is a Qwen3-0.6B decode step (28 transformer layers, + bf16). The trace captures ttnn ops: `EmbedDeviceIdsInto`, `sdpa_decode`, + `paged_update_cache`, RmsNorm, RoPE, etc. + +## The interleaving pattern (our decode driver) + +Per replay step, between `replay_mesh_trace` calls, the driver issues ~9 +host→device copies into persistent (pre-allocated, stable-address) buffers: + +- 4× rope cos/sin refresh (bf16) +- 2× RAC idx + page_table (INT32 ROW_MAJOR) — read by `paged_update_cache` +- 2× PA meta: page_table + cur_pos (INT32 ROW_MAJOR) — read by `sdpa_decode` +- 1× decode ids (UINT32) — read by the captured embedding + +Each copy is `ttnn::copy_to_device(host_tensor, persistent_device_tensor)` +which (per `ttnn/core/tensor/tensor_ops.cpp:182`) calls +`enqueue_write_tensor` → mesh CQ write path → `device_dispatch::write_to_core` +(`fd_mesh_command_queue.cpp:579`). + +## Bisection (the toxic class is idx/PA-meta INT32 copies) + +Qwen3-0.6B, "Hello", `--max-tokens 80`, `VT_TT_HOST_FREE_DECODE=1`. + +We have env-gated skip flags that suppress specific copy classes after the +first capture (stale device content — mechanics test only, numerically wrong): + +| skip flag(s) ON | copies still interleaved | result | replays before hang | +|---|---|---|---| +| (none, `VT_TT_RECAPTURE_EVERY=32`) | rope+idx+PAmeta+ids | HANG | 75 (32+32+11, gen3) | +| `VT_TT_NO_ROPE_REFRESH` | idx+PAmeta+ids | HANG | 39 | +| `VT_TT_NO_IDX_WARM` | rope+ids | **PASS** | 79 (RC=0, completed) | +| `VT_TT_NO_IDS_WARM` | rope+idx+PAmeta | HANG | 50 | +| all three (rope+idx+ids) | (none) | **PASS** | 79 (RC=0, completed) | + +**Conclusion:** the toxic copy class is the idx/PA-meta refresh +(`WarmRacIdx` + `WarmPaMeta`): `copy_to_device` into the INT32 +`page_table`/`update_idxs`/`cur_pos` buffers the captured +`sdpa_decode`/`paged_update_cache` reads. With idx/PA-meta OFF (rope+ids +ON), 79 replays complete cleanly; with idx/PA-meta ON it hangs regardless of +the others. The wall-count variance (39/50/75) is run-to-run timing noise, +not evidence that rope/ids contribute (the gate is binary on idx/PA-meta). + +## Upstream mesh-trace machinery is NOT the cause + +I ported the upstream `tests/tt_metal/distributed/test_mesh_trace.cpp` +`Sanity` test plus four scratch variants onto this checkout on the same card. +ALL PASS: + +- `MeshTraceTestSuite.Sanity`: 10 capture gens × 4 traces × 40 replays + + release — PASS (21 s). +- `ScratchConsecutiveReplayOneTrace`: 1 program, 120 consecutive replays of + ONE trace id, no interleaved traffic — PASS. +- `ScratchConsecutiveReplayFinishEach`: same + `Finish()` after every replay + (mimics our per-step sync) — PASS. +- `ScratchRecaptureGenerations`: 10 gens × 16 replays + release — PASS. +- `ScratchBigTraceReplay`: 50-program trace (~50 MB), 120 replays — PASS (47 s). + +The raw `replay_mesh_trace` / `release_trace` / `MeshTraceBuffer` lifetime is +sound. The wall only appears under our interleaving of `copy_to_device` into +trace-input buffers between replays. + +## Mechanism (code read) + +1. `copy_to_device` (`ttnn/core/tensor/tensor_ops.cpp:182`) → + `enqueue_write_tensor` → mesh CQ `write_to_core` + (`tt_metal/distributed/fd_mesh_command_queue.cpp:579`). +2. `write_to_core` and the interleaved-buffer write dispatch params both take + `expected_num_workers_completed` + (`tt_metal/impl/buffers/dispatch.cpp:486, 509, 520, 737, 985, 1304, 1350`). +3. The trace replay path `update_worker_state_post_trace_execution` + (`tt_metal/impl/trace/dispatch.cpp:31-72`) assigns + `expected_num_workers_completed` from the trace descriptor and calls + `config_buffer_mgr.mark_completely_full(expected_num_workers_completed[i])` + per replay (`dispatch.cpp:223`). +4. The trace's embedded `DISPATCH_WAIT(count=expected_workers)` cmds assume + the count starts from a known state at each replay (the reset mechanism + described in #7793). Interleaved `write_to_core` between replays perturbs + this counter; after ~38–50 such perturbations the next replay's + `DISPATCH_WAIT` never completes. + +Writes during *capture* are correctly forbidden +(`TT_FATAL(!trace_id_.has_value(), "Writes are not supported during trace +capture.")` at `fd_mesh_command_queue.cpp:572`), but writes between *replays* +(trace live, not capturing) are allowed and are the trigger. + +## Related + +- #7793 — the within-trace version of this bug (the `DISPATCH_WAIT(count)` / + `expected_num_workers_completed` / `num_completion_worker_cores` family). +- #7978 — fix for the within-trace case. +- #19248 — `expected_num_workers_completed` counter-wrap hangs in MeshCQ. +- #17696 — non-deterministic MeshDevice trace replay hangs (blocking + behavior). + +## Ask + +1. Is there a sanctioned way to issue host→device writes into trace-input + buffers between replays without perturbing + `expected_num_workers_completed`? (e.g. a write path that bypasses the + dispatch worker-completion accounting, or an explicit counter-reset / + `Finish` protocol between replays.) +2. If not, should the between-replays case get the same treatment #7978 gave + the within-trace case — i.e. a count-reset protocol so interleaved writes + don't desync the next replay's `DISPATCH_WAIT`? + +## Workaround (in our driver) + +`VT_TT_RECAPTURE_EVERY=N`: destroy + re-capture the graph every N replays. +N=32 gives ~2 healthy capture generations (~75 replays) before the same wall +reappears on the re-captured trace — so the leak also accumulates across +capture cycles, not just within one. Not a viable production fix. + +The long-term direction for our driver is to move per-step state advancement +(cur_pos, page_table, update_idx) inside the captured graph (the upstream +`executor.py` traced-decode pattern: 0–1 copies/step), eliminating the +interleaved writes entirely. But the between-replays write-into-trace-input +case seems like a genuine dispatch-accounting gap worth flagging. + +## Repro + +**Honest status: the trigger does NOT reproduce in any minimal standalone test +I tried.** I can only reproduce it in the full vllm.cpp decode driver. + +Minimal repros that all PASS (no hang): +- C++ scratch test in `test_mesh_trace.cpp` (`ScratchInterleavedWriteBetweenReplays`): + capture an eltwise-binary program, replay 120× with `EnqueueWriteMeshBuffer` + into its DRAM src buffer between replays — PASS. +- Python script (`repro_trace_replay_write_hang.py`): capture a single + `ttnn.transformer.paged_scaled_dot_product_attention_decode` (the exact op + our driver uses) with persistent device page_table/cur_pos, replay 120× with + `ttnn.copy` into those INT32 buffers between replays (the exact toxic class) + — PASS. With `blocking=False` + readback sync (matching the driver) — PASS. +- Same script with NUM_OPS=8 and NUM_OPS=28 (28 chained decode calls in one + trace, matching our layer count) — PASS. +- Adding `ttnn.experimental.paged_update_cache` (the RAC op, with the exact + height-sharded L1 input `[1,1,32,128]` our driver uses) to the captured + trace, and refreshing page_table + cur_pos + update_idxs (all three toxic + buffers) between replays — PASS (both 1× and 28×). + +So the trigger is not the op, not the op count, not the RAC op, not the +sharded input, not the ttnn wrapper, not the INT32 write class in isolation, +not blocking=False. It only manifests in the full driver trace, which +differs from the repros in: (1) a diverse op mix (embedding + 28× +[RmsNorm+RoPE+sdpa_decode+paged_update_cache+...] with differing +`num_completion_worker_cores` per op, vs identical chained RAC+sdpa pairs); +(2) our persistent-tensor pool + `from_vector`/`copy_to_device` host-staging +path (vs the repro's `as_tensor` fresh device tensor + `ttnn.copy`). + +Filing now with the honest non-minimal framing, in case the between-replays +write case is a known limitation with a sanctioned workaround, and to ask +whether `ttnn::copy_to_device`-into-trace-input-between-replays is expected +to be safe. I can share the full driver reproducer if useful; porting it to +a self-contained tt-metal test would require reconstructing the decode-layer +op mix. 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/BENCHMARKS.md b/docs/BENCHMARKS.md index cc516c4e5..0a7606c3d 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -507,6 +507,7 @@ built on it rather than keeping the flattering one. | ROCm `d=128` decode arm (`BACKEND-ROCM`, [#382](https://github.com/mudler/vllm.cpp/issues/382)) | **DIRECTIONAL, not binding.** gfx1200, both sides in the pinned oracle container, 1024/128 c1, 8 prompts, 3 reps: TPOT 42.40 -> **11.67 ms** with `VT_ATTN_DECODE_D128=1`; vLLM `555967922` 6.68 ms, so **6.35x -> 1.75x** | Harnesses differ (oracle over HTTP, ours in-process) and no same-tool per-call trace exists, so [#488](https://github.com/mudler/vllm.cpp/issues/488) stays open. Owed: decode-windowed `rocprofv3` both sides | | Tenstorrent Blackhole (`BACKEND-TENSTORRENT`) | **NOT APPLICABLE (speed).** Correctness: OPT-125m STRICT 6/6 e2e on real hardware. Qwen3-0.6B has a device-specific golden and short 4-token warm smoke (~0.28 tok/s), not a completed speed run | Full 16x16 Qwen3 gate, then device-resident tensors + `ttnn::sdpa_decode` before any performance comparison. [Spec](../.agents/specs/tenstorrent-backend.md) | | Mistral-7B-v0.3 on Tenstorrent (`BACKEND-TENSTORRENT-MISTRAL`) | **PENDING (speed).** First data point: 4.26 tok/s warm, batch 1, 32 tok, single run on a P150. Not a gate, not reproduced. No vLLM ratio exists or can (no TT backend). Correctness 16/16 | Reproduce idle with a same-binary A/B before quoting. [Record](../.agents/benchmark-record.md), [spec](../.agents/specs/tenstorrent-mistral.md) | +| Host-free decode graph (`BACKEND-TENSTORRENT-HOST-FREE-FORWARD`) | **PENDING (operator gate).** Implementer P150 run (2026-08-16): Qwen3-0.6B, 80 tok, 5.8x vs eager, 79 replays, 22/22 vs copy baseline. Env-gated. No vLLM ratio | Operator rerun, then full-engine `test_qwen3_paged_engine`. [Spec](../.agents/specs/tenstorrent-host-free-forward.md) | | Prompt logprobs (`SAMPLE-PROMPT-LOGPROBS`, #223) | **NO number measured, claimed or owed.** Correctness-only, CPU. Upstream ships this path explicitly unoptimized (`gpu_model_runner.py:5622-5623`); a step where no request asks is unchanged | Floor if one is ever wanted: vLLM's own `prompt_logprobs=k`, same model and prompt | | `logprobs_mode` (`SAMPLE-LOGPROB-TOKEN-IDS`, #238) | **NO number measured, claimed or owed.** Correctness-only, CPU. One [n, vocab] device->host copy per step when a processed mode is engaged, nothing when not | Nothing to close: observation modes, not a path vLLM optimizes either | | `logprob_token_ids` scoring (#264) | **No number owed:** correctness-only, CPU-gated; inert unless set | Owed once the OpenAI field is wired | diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 34b8238c1..d18bd64fa 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` | `0` (off) | Seconds between SSE comment keepalives (`:\n\n`) on silent streams; `<=0` disables. OFF is the default because vLLM's streaming endpoints emit no comment frame at all, and some SSE clients — including vLLM's own `vllm bench serve` — cannot resynchronise after one, silently counting the request FAILED while the server completes it. Set a positive value (clamped to 600) only behind a proxy whose inactivity timeout would otherwise drop a long prefill, and expect strict clients to break. See [#931](https://github.com/mudler/vllm.cpp/issues/931) | | `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: residual-RMS + RoPE stay on device at decode shapes, the dense decode graph captures, and `cur_pos` advances on-device inside the trace. Default decode is unchanged without it. An implementer P150 run completed 79 replays; the operator gate and the full-engine golden are still owed A second concurrent request (a padded batch-size change) REFUSES under this flag rather than returning wrong tokens; see #1105 | | `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 | @@ -187,6 +188,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/docs/FEATURES.md b/docs/FEATURES.md index ae04cee39..074bf8fa9 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -265,6 +265,7 @@ both refuse, naming what is missing. | ROCm | W0 verified on 5 gfx archs; dense and GDN models run all-native. Strict CPU parity is open in the measured near-tie regime (#269) | 44 registered ops including full GDN; ctest-green gfx1151/1103/1100/1201/1200 ([#41](https://github.com/mudler/vllm.cpp/issues/41)). APU managed allocation is unverified. [ROCM.md](ROCM.md) | ✅ | ✅ | | XPU / TPU | ☐ | ✅ | ◐ | ☐ | | Tenstorrent Blackhole | ◐ `ACTIVE`, OPT-125m 6/6; Qwen3-0.6B wired; Mistral-7B-v0.3 16/16 on P150 ([spec](../.agents/specs/tenstorrent-mistral.md)). 16x16 rerun and residual-RMS owed ([spec](../.agents/specs/tenstorrent-backend.md)) | ✅ | ☐ | ☐ | +| Tenstorrent host-free decode | ◐ env-gated `VT_TT_HOST_FREE_DECODE`; implementer P150 79-replay/5.8x. Default inert. New batch after capture refused. Engine golden owed | ☐ | ☐ | ☐ | CUDA runtime-verified on GB10 (sm_121a), Jetson Thor (sm_110) and Jetson AGX Orin (sm_87). sm_110 has no CUTLASS FP4 tensor-core kernels and no `fp4-mma`, diff --git a/docs/STATUS.md b/docs/STATUS.md index 7d46436ef..a589bb4c4 100644 --- a/docs/STATUS.md +++ b/docs/STATUS.md @@ -1622,6 +1622,10 @@ Gemma4/ROCm env split: public `VT_GEMMA4_EXPERT_VRAM_MB` caps expert LRU in posi `BACKEND-TENSTORRENT-MISTRAL`: `ACTIVE`: Mistral-7B-v0.3 gated on a Blackhole P150, 16/16 prompts (12/16 strict token-exact, 4/16 inside the near-tie band, 0 forward-divergent), max gap 0.062 nats. `MistralForCausalLM` is allowlisted by exact match, so `Mistral3ForConditionalGeneration` (#387, unported) still falls through. Correctness only -- no speed claim. +`BACKEND-TENSTORRENT-TRACE-RUNNER`: `SPIKE`: NO-GO for T=1 decode capture. Measured `to_vector` abort. Decode capture moved to `BACKEND-TENSTORRENT-HOST-FREE-FORWARD`. Prefill capture still unaudited. + +`BACKEND-TENSTORRENT-HOST-FREE-FORWARD`: `ACTIVE`: env-gated `VT_TT_HOST_FREE_DECODE` decode-graph capture. Implementer P150 run of Qwen3-0.6B, 80 tokens: 79 replays, no hang, 5.8x vs eager, 22/22 vs the per-step-copy baseline. Default path inert. Operator gate and full-engine golden still owed. A new batch after the first capture is refused. + **Platform SELECTION is the one non-additive site, and is now gated.** A platform missing from `CurrentPlatform()`'s hardcoded walk registers and answers correctly but is NEVER selected, with no compiler diagnostic. `test_platform` diff --git a/scripts/check-agent-record.py b/scripts/check-agent-record.py index 026d10591..5ea2ce53b 100644 --- a/scripts/check-agent-record.py +++ b/scripts/check-agent-record.py @@ -299,14 +299,22 @@ # MistralForCausalLM on TT + device-aware SACRED gate. Reuses Qwen3-dense # forward; no new kernel. Pending 7B checkpoint + vLLM oracle for the e2e # gate. - # 83 since 2026-08-16: +`BACKEND-GATE-CUDA-LLAMACPP` (#979), the llama.cpp - # floor on a CURRENT CUDA card. Neither existing llama.cpp row covers it: - # `BACKEND-GATE-CPU-LLAMACPP` is the CPU floor and - # `BACKEND-GATE-CUDA-LLAMACPP-LEGACY` is scoped to Pascal/Volta/Turing, - # where vLLM has no entry at all. The four-way Qwen3.8-27B campaign needs - # it because llama.cpp is the ONLY comparator that runs the Q4_K_M arm: - # vLLM removed GGUF from its tree at our pin. INVENTORIED, no run. - "BACKEND": (AGENTS / "backend-matrix.md", 83), + # 83 since 2026-08-16: +`BACKEND-GATE-CUDA-LLAMACPP` (#979), the llama.cpp + # floor on a CURRENT CUDA card. Neither existing llama.cpp row covers it: + # `BACKEND-GATE-CPU-LLAMACPP` is the CPU floor and + # `BACKEND-GATE-CUDA-LLAMACPP-LEGACY` is scoped to Pascal/Volta/Turing, + # where vLLM has no entry at all. The four-way Qwen3.8-27B campaign needs + # it because llama.cpp is the ONLY comparator that runs the Q4_K_M arm: + # vLLM removed GGUF from its tree at our pin. INVENTORIED, no run. + # 84 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. + # 85 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", 85), } ENGINE_MATRIX = AGENTS / "engine-matrix.md" diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index 36cab9ceb..a54b1b3a5 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -173,6 +173,13 @@ VT_ROCM_HIPBLASLT VT_SILU_FP4_FAST VT_SPEC_TRACE VT_SWIZZLE_IN_QUANT +VT_TT_DUMP_KV +VT_TT_NO_IDS_WARM +VT_TT_NO_IDX_WARM +VT_TT_NO_ROPE_REFRESH +VT_TT_RECAPTURE_EVERY +VT_TT_SDPA_TEST +VT_TT_SHARD_Q VT_V4_DECODE_GRAPH VT_V4_DEVICE_ATTN VT_V4_DEVICE_GLUE diff --git a/scripts/probe_plus_one_capture.py b/scripts/probe_plus_one_capture.py new file mode 100644 index 000000000..14c569309 --- /dev/null +++ b/scripts/probe_plus_one_capture.py @@ -0,0 +1,31 @@ +# Spike probe: is ttnn.plus_one capture-safe + does it advance on replay? +import torch +import ttnn + +device = ttnn.open_mesh_device(l1_small_size=0, trace_region_size=50 * 1024 * 1024) +try: + dram = ttnn.DRAM_MEMORY_CONFIG + # Persistent INT32 cur_pos [1] on device, seeded to 5. + cur_pos = torch.tensor([5], dtype=torch.int32) + tt_cp = ttnn.as_tensor(cur_pos, device=device, dtype=ttnn.int32, memory_config=dram) + + # Warm plus_one (compile + program cache). + ttnn.plus_one(tt_cp, skip_negative_entries=True) + ttnn.synchronize_device(device) + print(f"[probe] after warm plus_one, cur_pos = {ttnn.to_torch(tt_cp).tolist()}", flush=True) + + # Capture: a no-op-ish body + plus_one at the end. + tid = ttnn.begin_trace_capture(device, cq_id=0) + ttnn.plus_one(tt_cp, skip_negative_entries=True) + ttnn.end_trace_capture(device, tid, cq_id=0) + + # Replay 5 times; each should increment cur_pos by 1. + for i in range(5): + ttnn.execute_trace(device, tid, cq_id=0, blocking=True) + v = ttnn.to_torch(tt_cp).tolist() + print(f"[probe] replay {i+1}: cur_pos = {v}", flush=True) + + ttnn.release_trace(device, tid) + print("[probe] DONE: plus_one is capture-safe and advances on replay", flush=True) +finally: + ttnn.close_mesh_device(device) diff --git a/scripts/repro_trace_replay_write_hang.py b/scripts/repro_trace_replay_write_hang.py new file mode 100644 index 000000000..f0d4cc5af --- /dev/null +++ b/scripts/repro_trace_replay_write_hang.py @@ -0,0 +1,120 @@ +# Scratch repro (vllm.cpp TT backend bisect, 2026-08-16) — NOT upstream. +# +# Reproduce the ~38-replay hang from the vllm.cpp decode driver as a standalone +# ttnn script. The driver's toxic pattern: capture a trace of a ttnn op whose +# INT32 inputs (page_table, cur_pos) are persistent device tensors, then replay +# N times with ttnn.copy (copy_to_device) into those inputs between replays, +# trace staying live. Mirrors tests/ttnn/unit_tests/operations/sdpa/sdpa_test_utils.py +# run_test_paged_attention_trace path for the op setup. +# +# Usage: +# source /home/lu_zero/Sources/tt/env-tt-common.sh +# python repro_trace_replay_write_hang.py +import os +import torch +import ttnn + + +def to_paged_cache(cache, batch, num_kv, max_num_blocks_per_seq, block_size, head_dim): + return ( + cache.reshape(batch, num_kv, max_num_blocks_per_seq, block_size, head_dim) + .transpose(1, 2) + .reshape(batch * max_num_blocks_per_seq, num_kv, block_size, head_dim) + ) + + +def main(): + NUM_ITERS = int(os.environ.get("REPRO_ITERS", "120")) + DO_WRITE = os.environ.get("REPRO_NO_WRITE") != "1" + + device = ttnn.open_mesh_device(l1_small_size=0, trace_region_size=50 * 1024 * 1024) + try: + b, nh, nkv, s, d = 1, 8, 1, 512, 128 + block_size = 32 + max_num_blocks_per_seq = s // block_size + max_num_blocks = b * s // block_size + scale = d**-0.5 + dram = ttnn.DRAM_MEMORY_CONFIG + + K = torch.randn(b, nkv, s, d, dtype=torch.bfloat16) + V = torch.randn(b, nkv, s, d, dtype=torch.bfloat16) + paged_k = to_paged_cache(K, b, nkv, max_num_blocks_per_seq, block_size, d) + paged_v = to_paged_cache(V, b, nkv, max_num_blocks_per_seq, block_size, d) + permutation = torch.randperm(max_num_blocks) + reverse_permutation = torch.argsort(permutation) + page_table = reverse_permutation.reshape(b, max_num_blocks_per_seq).to(torch.int32) + paged_k_shuffled = paged_k[permutation] + paged_v_shuffled = paged_v[permutation] + + tt_K = ttnn.as_tensor(paged_k_shuffled, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, memory_config=dram) + tt_V = ttnn.as_tensor(paged_v_shuffled, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, memory_config=dram) + tt_page_table = ttnn.as_tensor(page_table, device=device, dtype=ttnn.int32, memory_config=dram) + Q = torch.randn(1, b, nh, d, dtype=torch.bfloat16) + tt_Q = ttnn.as_tensor(Q, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, memory_config=dram) + start_indices = torch.tensor([s - 1], dtype=torch.int32) + tt_cur_pos = ttnn.as_tensor(start_indices, device=device, dtype=ttnn.int32, memory_config=dram) + + NUM_OPS = int(os.environ.get("REPRO_NUM_OPS", "1")) # sdpa calls chained per trace + USE_RAC = os.environ.get("REPRO_NO_RAC") != "1" # include paged_update_cache in the trace + + # paged_update_cache needs a height-sharded K/V input, shard [nkv_pad, d] + # on one core (nkv padded to TILE_HEIGHT=32). Matches our driver's RAC. + nkv_pad = 32 + shard_grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(0, 0))}) + shard_spec = ttnn.ShardSpec(shard_grid, [nkv_pad, d], ttnn.ShardOrientation.ROW_MAJOR) + k_input_memcfg = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, shard_spec) + k_input = torch.zeros(1, 1, nkv_pad, d, dtype=torch.bfloat16) + tt_k_input = ttnn.as_tensor(k_input, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, memory_config=k_input_memcfg) + v_input = torch.zeros(1, 1, nkv_pad, d, dtype=torch.bfloat16) + tt_v_input = ttnn.as_tensor(v_input, device=device, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, memory_config=k_input_memcfg) + update_idxs = [s - 1] + tt_update_idxs = ttnn.as_tensor(torch.tensor(update_idxs, dtype=torch.int32), device=device, dtype=ttnn.int32, memory_config=dram) + + # Warm both ops (compile + program cache). + if USE_RAC: + ttnn.experimental.paged_update_cache(tt_K, tt_k_input, update_idxs=[], update_idxs_tensor=tt_update_idxs, page_table=tt_page_table) + ttnn.experimental.paged_update_cache(tt_V, tt_v_input, update_idxs=[], update_idxs_tensor=tt_update_idxs, page_table=tt_page_table) + _ = ttnn.transformer.paged_scaled_dot_product_attention_decode( + tt_Q, tt_K, tt_V, tt_page_table, is_causal=True, cur_pos_tensor=tt_cur_pos, scale=scale) + ttnn.synchronize_device(device) + + # Capture: paged_update_cache (RAC) + paged sdpa decode, chained NUM_OPS times. + tid = ttnn.begin_trace_capture(device, cq_id=0) + tt_out = tt_Q + for _ in range(NUM_OPS): + if USE_RAC: + ttnn.experimental.paged_update_cache(tt_K, tt_k_input, update_idxs=[], update_idxs_tensor=tt_update_idxs, page_table=tt_page_table) + ttnn.experimental.paged_update_cache(tt_V, tt_v_input, update_idxs=[], update_idxs_tensor=tt_update_idxs, page_table=tt_page_table) + tt_out = ttnn.transformer.paged_scaled_dot_product_attention_decode( + tt_out, tt_K, tt_V, tt_page_table, is_causal=True, cur_pos_tensor=tt_cur_pos, scale=scale) + ttnn.end_trace_capture(device, tid, cq_id=0) + + print(f"[repro] captured; now {NUM_ITERS} replays with do_write={DO_WRITE}", flush=True) + for i in range(NUM_ITERS): + if DO_WRITE: + # Refresh the INT32 trace-input buffers in place (the toxic + # pattern: page_table + cur_pos + update_idxs, matching + # WarmPaMeta + WarmRacIdx). + page_table[0, 0] = i % max_num_blocks_per_seq + tt_pt_fresh = ttnn.as_tensor(page_table, device=device, dtype=ttnn.int32, memory_config=dram) + ttnn.copy(tt_pt_fresh, tt_page_table) + cp = torch.tensor([s - 1], dtype=torch.int32) + tt_cp_fresh = ttnn.as_tensor(cp, device=device, dtype=ttnn.int32, memory_config=dram) + ttnn.copy(tt_cp_fresh, tt_cur_pos) + ui = torch.tensor([s - 1], dtype=torch.int32) + tt_ui_fresh = ttnn.as_tensor(ui, device=device, dtype=ttnn.int32, memory_config=dram) + ttnn.copy(tt_ui_fresh, tt_update_idxs) + # Match the driver exactly: blocking=False, NO synchronize_device, + # the .cpu() readback is the sync point. + ttnn.execute_trace(device, tid, cq_id=0, blocking=False) + if i % 5 == 0: + print(f"[repro] iter {i}/{NUM_ITERS} ok", flush=True) + _ = tt_out.cpu() + print(f"[repro] DONE: {NUM_ITERS} replays, no hang", flush=True) + ttnn.release_trace(device, tid) + finally: + ttnn.close_mesh_device(device) + + +if __name__ == "__main__": + main() diff --git a/src/vllm/model_executor/models/qwen3.cpp b/src/vllm/model_executor/models/qwen3.cpp index 614356f7c..80ed8b3b8 100644 --- a/src/vllm/model_executor/models/qwen3.cpp +++ b/src/vllm/model_executor/models/qwen3.cpp @@ -32,6 +32,7 @@ // overlaps the host-side alloc syncs with GPU compute; it is kept as byte-safe // hygiene + code sharing, not a measured TTFT lever. The real dense-TTFT lever // is the RoPE cos|sin cache below. +#include #include "vllm/model_executor/models/qwen3.h" #include @@ -55,6 +56,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 { @@ -77,6 +79,13 @@ using v1::CommonAttentionMetadata; // so the Qwen3-dense (0.6B/4B) forward is byte-identical (same vt:: op order). using namespace dense_attn; +// TT-only dump. The getenv is paid only on kTENSTORRENT so a CUDA replay +// step does not walk the environment twice. +bool TtDumpKv(const Dev& d) { + return d.q.device.type == vt::DeviceType::kTENSTORRENT && + std::getenv("VT_TT_DUMP_KV") != nullptr; +} + // Dense SwiGLU MLP (qwen3.py::Qwen3MLP=Qwen2MLP): merged gate_up_proj -> // SiluAndMul -> down_proj. `dh2` is the post-norm hidden [T,H] bf16. // @@ -456,6 +465,14 @@ std::vector Qwen3DenseModel::Forward( const int64_t n_out = dlogits.t().shape[0]; std::vector logits(static_cast(n_out) * config.vocab_size); dlogits.Download(d, logits.data()); + if (TtDumpKv(d)) { + int argmax = 0; + for (int64_t i = 1; i < n_out * config.vocab_size; ++i) + if (logits[static_cast(i)] > logits[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] Forward eager argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits[0], logits[1], logits[2], logits[3], logits[4]); + } return logits; } @@ -468,6 +485,16 @@ ForwardLogits Qwen3DenseModel::ForwardDevice( DBuf dlogits = ForwardBody(d, token_ids, positions, attn_meta, attn_kv, weights, config, logits_indices); const int64_t n_out = dlogits.t().shape[0]; + if (TtDumpKv(d)) { + std::vector logits_dump(static_cast(n_out * config.vocab_size)); + dlogits.Download(d, logits_dump.data()); + int argmax = 0; + for (size_t i = 1; i < logits_dump.size(); ++i) + if (logits_dump[i] > logits_dump[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] ForwardDevice eager argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits_dump[0], logits_dump[1], logits_dump[2], logits_dump[3], logits_dump[4]); + } return WrapDeviceLogits(d, std::move(dlogits), n_out, config.vocab_size); } @@ -561,11 +588,19 @@ struct Qwen3DenseDecodeGraph::Impl { b.SupportsGraphCapture(); } ~Impl() { - if (std::getenv("VT_DECODE_GRAPH_STATS") != nullptr) + if (std::getenv("VT_DECODE_GRAPH_STATS") != nullptr) { + std::string extra; + if (replay_steps > 0) { + extra = "; replay branch avg " + + std::to_string(static_cast(replay_ns) / 1e6 / + static_cast(replay_steps)) + + " ms/step over " + std::to_string(replay_steps) + " steps"; + } std::fprintf(stderr, "[Qwen3DenseDecodeGraph] dense decode graph: %lld total replays " - "across %zu captured size(s)\n", - static_cast(replays), slots.size()); + "across %zu captured size(s)%s\n", + static_cast(replays), slots.size(), extra.c_str()); + } Backend& b = vt::GetBackend(queue.device.type); for (auto& kv : slots) if (kv.second.graph != nullptr) b.DestroyGraph(kv.second.graph); @@ -615,6 +650,10 @@ struct Qwen3DenseDecodeGraph::Impl { std::map slots; // padded size S -> slot int64_t replays = 0; // total replays (diagnostics) bool any_captured = false; // diagnostics: at least one live graph + // Steady-state timing (VT_DECODE_GRAPH_STATS): wall time of the replay + // branch (warm copies + ReplayGraph; excludes the caller's logits readback). + int64_t replay_ns = 0; + int64_t replay_steps = 0; }; Qwen3DenseDecodeGraph::Qwen3DenseDecodeGraph(const Qwen3DenseWeights& weights, @@ -644,6 +683,16 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( if (!impl_->enabled || S < 0) { DBuf lg = ForwardBody(d, token_ids, positions, attn_meta, attn_kv, impl_->weights, impl_->config, kNoGather); + if (TtDumpKv(d)) { + std::vector logits_dump(static_cast(vocab)); + lg.Download(d, logits_dump.data()); + int argmax = 0; + for (int64_t i = 1; i < vocab; ++i) + if (logits_dump[static_cast(i)] > logits_dump[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] eager path argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits_dump[0], logits_dump[1], logits_dump[2], logits_dump[3], logits_dump[4]); + } return WrapDeviceLogits(d, std::move(lg), B, vocab); } @@ -655,11 +704,81 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( CommonAttentionMetadata pam; BuildPaddedDecodeAttn(S, token_ids, positions, attn_meta, ptok, ppos, pam); + // Debug: dump attention metadata for comparison + if (TtDumpKv(d)) { + fprintf(stderr, "[TT-DUMP-META] S=%lld B=%lld\n", (long long)S, (long long)B); + fprintf(stderr, "[TT-DUMP-META] real: num_reqs=%d num_tokens=%d slot0=%lld seq_len0=%d bt_cols=%d\n", + attn_meta.num_reqs, attn_meta.num_actual_tokens, + attn_meta.slot_mapping.empty() ? -1LL : (long long)attn_meta.slot_mapping[0], + attn_meta.seq_lens.empty() ? -1 : attn_meta.seq_lens[0], + attn_meta.block_table_num_cols); + fprintf(stderr, "[TT-DUMP-META] pad: num_reqs=%d num_tokens=%d slot0=%lld seq_len0=%d bt_cols=%d\n", + pam.num_reqs, pam.num_actual_tokens, + pam.slot_mapping.empty() ? -1LL : (long long)pam.slot_mapping[0], + pam.seq_lens.empty() ? -1 : pam.seq_lens[0], + pam.block_table_num_cols); + fprintf(stderr, "[TT-DUMP-META] real pos0=%d pad pos0=%d\n", + positions.empty() ? -1 : positions[0], + ppos.empty() ? -1 : ppos[0]); + } + // A block-table column-count change reallocates the persistent block_table (the // captured H2D copy's source address moves) -> invalidate this slot's graph and // 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); + } + // R2: seed the on-device-advanced cur_pos BEFORE WarmRacIdx, so the RAC + // path can alias update_idxs to it (eliminating the per-replay + // update_idxs copy_to_device — the toxic ~38-replay hang class). + if (!pam.seq_lens.empty()) { + vt::tenstorrent::WarmDecodePos( + pam.seq_lens.data(), static_cast(pam.num_reqs)); + } + 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, + pam.block_table_tensor.data(), + static_cast(pam.block_table_num_cols), + pam.seq_lens.data()); + // 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); @@ -668,13 +787,67 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( s.warm = false; } - // Fast path: this size's graph is captured. Embed OUTSIDE the graph into the - // persistent hidden buffer, then relaunch the captured layer region. - if (s.captured) { - EmbedInto(d, *s.hidden, s.token_ids, impl_->weights, impl_->config); + // Fast path: this size's graph is captured. On TT, refresh the persistent + // decode-ids tensor (allocation-free) and replay — the embedding itself is + // INSIDE the captured region, so a replay step performs zero eager device + // allocations (eager alloc/free churn around a live trace hung the device + // ~60 replays in). CUDA keeps the outside-the-graph EmbedInto. + // VT_TT_RECAPTURE_EVERY=N (TT only): destroy and re-capture the graph every + // N replays. WORKAROUND for the deterministic ~38-replay completion hang on + // this tt-metal build: a replayed mesh trace stops completing (futex wait + // in the post-replay readback) after ~38 replays of one trace id, + // independent of interleaved eager-copy count and readback count. + // Re-capturing resets the per-trace device state; the eager re-warm step + // and capture step run with NO live trace (DestroyGraph releases it), so + // eager allocations are legal there. + bool do_replay = s.captured; + if (do_replay && d.q.device.type == vt::DeviceType::kTENSTORRENT) { + const char* rc_env = std::getenv("VT_TT_RECAPTURE_EVERY"); + const int rc_n = rc_env != nullptr ? std::atoi(rc_env) : 0; + if (rc_n > 0 && s.replays >= rc_n) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, + "[TT-STEP] recapture: destroying graph after %lld replays " + "(every %d)\n", + static_cast(s.replays), rc_n); + b.DestroyGraph(s.graph); + s.graph = nullptr; + s.captured = false; + s.warm = false; + do_replay = false; + } + } + if (do_replay) { + const auto replay_t0 = std::chrono::steady_clock::now(); + if (d.q.device.type == vt::DeviceType::kTENSTORRENT) { + vt::tenstorrent::WarmDecodeIds( + s.token_ids.data(), static_cast(s.token_ids.size())); + } else { + EmbedInto(d, *s.hidden, s.token_ids, impl_->weights, impl_->config); + } + if (TtDumpKv(d)) { + std::vector pre(static_cast(vocab)); + s.logits->Download(d, pre.data()); + fprintf(stderr, "[TT-DUMP-LOGITS] pre-replay first5=[%f,%f,%f,%f,%f]\n", + pre[0], pre[1], pre[2], pre[3], pre[4]); + } b.ReplayGraph(impl_->queue, s.graph); ++s.replays; ++impl_->replays; + impl_->replay_ns += std::chrono::duration_cast( + std::chrono::steady_clock::now() - replay_t0) + .count(); + ++impl_->replay_steps; + if (TtDumpKv(d)) { + std::vector logits_dump(static_cast(vocab)); + s.logits->Download(d, logits_dump.data()); + int argmax = 0; + for (int64_t i = 1; i < vocab; ++i) + if (logits_dump[static_cast(i)] > logits_dump[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] replay step argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits_dump[0], logits_dump[1], logits_dump[2], logits_dump[3], logits_dump[4]); + } return ViewDeviceLogits(s.logits->ptr(), d.q.device, B, vocab); } @@ -682,10 +855,36 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( // this size by the previous (eager) step. CAPTURE the layer region once, // instantiate the graph, then launch it. if (s.warm) { - EmbedInto(d, *s.hidden, s.token_ids, impl_->weights, impl_->config); + const bool tt_dev = d.q.device.type == vt::DeviceType::kTENSTORRENT; + if (tt_dev) { + // Stage ids for the captured embedding (outside capture). + vt::tenstorrent::WarmDecodeIds( + s.token_ids.data(), static_cast(s.token_ids.size())); + } else { + EmbedInto(d, *s.hidden, s.token_ids, impl_->weights, impl_->config); + } b.BeginCapture(impl_->queue); + if (tt_dev) { + // Capture-safe embedding over the persistent ids tensor, writing the + // persistent hidden shadow the layer region reads. + Tensor dtab = ResidentWeight(d, impl_->weights.embed_tokens, + {impl_->config.vocab_size, + impl_->config.hidden_size}); + vt::tenstorrent::EmbedDeviceIdsInto( + s.hidden->ptr(), S, impl_->config.hidden_size, dtab.data, + impl_->config.vocab_size, impl_->config.hidden_size, + static_cast(s.token_ids.size())); + } DBuf lg = ForwardLayers(d, s.hidden->t(), s.positions, s.attn_meta, attn_kv, impl_->weights, impl_->config, kNoGather); + // R2: advance cur_pos on-device (plus_one) INSIDE the captured trace, at + // the END of the body (after all reads of cur_pos in sdpa_decode/RAC). + // The NEXT replay sees cur_pos+1 — eliminating the per-replay cur_pos / + // update_idxs copy_to_device (the toxic ~38-replay hang class). + if (tt_dev && !pam.seq_lens.empty()) { + vt::tenstorrent::CaptureDecodePosAdvance( + static_cast(pam.num_reqs)); + } s.graph = b.EndCaptureGraph(impl_->queue); s.logits = std::make_unique(std::move(lg)); s.captured = true; @@ -698,6 +897,16 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( b.ReplayGraph(impl_->queue, s.graph); s.replays = 1; ++impl_->replays; + if (TtDumpKv(d)) { + std::vector logits_dump(static_cast(vocab)); + s.logits->Download(d, logits_dump.data()); + int argmax = 0; + for (int64_t i = 1; i < vocab; ++i) + if (logits_dump[static_cast(i)] > logits_dump[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] capture step argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits_dump[0], logits_dump[1], logits_dump[2], logits_dump[3], logits_dump[4]); + } return ViewDeviceLogits(s.logits->ptr(), d.q.device, B, vocab); } @@ -708,6 +917,17 @@ ForwardLogits Qwen3DenseDecodeGraph::Step( EmbedInto(d, *s.hidden, s.token_ids, impl_->weights, impl_->config); DBuf lg = ForwardLayers(d, s.hidden->t(), s.positions, s.attn_meta, attn_kv, impl_->weights, impl_->config, kNoGather); + // Debug: dump logits first 5 values + if (TtDumpKv(d)) { + std::vector logits_dump(static_cast(vocab)); + lg.Download(d, logits_dump.data()); + int argmax = 0; + for (int64_t i = 1; i < vocab; ++i) + if (logits_dump[static_cast(i)] > logits_dump[static_cast(argmax)]) + argmax = static_cast(i); + fprintf(stderr, "[TT-DUMP-LOGITS] cold step argmax=%d first5=[%f,%f,%f,%f,%f]\n", + argmax, logits_dump[0], logits_dump[1], logits_dump[2], logits_dump[3], logits_dump[4]); + } s.warm = true; s.captured = false; // lg is [S,vocab]; hand ownership out but expose only the first B (real) rows. 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..298a83dd0 100644 --- a/src/vt/tenstorrent/tenstorrent_device.h +++ b/src/vt/tenstorrent/tenstorrent_device.h @@ -3,6 +3,7 @@ // tenstorrent-backend.md). vllm.cpp original; no upstream mirror. #pragma once +#include #include // ttnn::MeshDevice (ttnn/api/ttnn/device.hpp) is a `using` alias for this real @@ -49,6 +50,127 @@ 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). +#ifdef VLLM_CPP_TENSTORRENT +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); +#else +inline void WarmPagedKvShadow(void*, void*, int64_t, int64_t, int64_t, int64_t, + int64_t) {} +#endif + +// 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). page_table is +// the block table (virtual→physical block mapping, int32, [num_reqs, cols]). +// positions are the sequence positions for this step (int32). Both are used to +// build the update_idxs (positions) and page_table tensors paged_update_cache +// reads at replay time. +#ifdef VLLM_CPP_TENSTORRENT +void WarmRacIdx(const void* slot_mapping_owner, const int64_t* slots, + int64_t num_slots, int64_t block_size, + const int32_t* page_table, int64_t page_table_cols, + const int32_t* positions); +#else +inline void WarmRacIdx(const void*, const int64_t*, int64_t, int64_t, + const int32_t*, int64_t, const int32_t*) {} +#endif + +// ITEM 5 (PA): warm persistent page_table + cur_pos device tensors. +// When `advance_on_device` is true, cur_pos is seeded once and then advanced +// in-trace via CaptureDecodePosAdvance (ttnn::plus_one); WarmPaMeta's cur_pos +// copy_to_device is skipped in steady-state replay (R2 on-device state advance). +#ifdef VLLM_CPP_TENSTORRENT +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); +#else +inline void WarmPaMeta(const int32_t*, int64_t, int64_t, int64_t, int64_t, + const int32_t*) {} +#endif + +// R2 (on-device state advance): seed the persistent cur_pos device tensor +// (= seq_lens - 1) for this step. Called on the capture/warm step (re-seed), +// NOT every replay step — the captured plus_one advances it on-device. +// Also warms the plus_one program (program cache) so CaptureDecodePosAdvance +// can run inside the trace without a "load new binaries during capture" fatal. +#ifdef VLLM_CPP_TENSTORRENT +void WarmDecodePos(const int32_t* seq_lens, int64_t num_reqs); +#else +inline void WarmDecodePos(const int32_t*, int64_t) {} +#endif + +// R2: capture ttnn::plus_one(cur_pos) at the END of the trace body (after all +// reads of cur_pos in sdpa_decode / paged_update_cache), so the NEXT replay +// sees cur_pos+1. Must be called INSIDE BeginCapture/EndCapture. +#ifdef VLLM_CPP_TENSTORRENT +void CaptureDecodePosAdvance(int64_t num_reqs); +#else +inline void CaptureDecodePosAdvance(int64_t) {} +#endif + +// HOST-FREE-DECODE: stage the persistent device decode-ids tensor (UINT32 [n], +// ROW_MAJOR) for THIS step's token ids, outside capture. The captured +// embedding reads this stable address; each replay step only refreshes its +// content (copy_to_device — allocation-free). No-op unless +// VT_TT_HOST_FREE_DECODE. +#ifdef VLLM_CPP_TENSTORRENT +void WarmDecodeIds(const int32_t* ids, int64_t n); +#else +inline void WarmDecodeIds(const int32_t*, int64_t) {} +#endif + +// HOST-FREE-DECODE: capture-safe embedding over the persistent decode-ids +// tensor (see WarmDecodeIds) into the hidden buffer whose host base is +// `out_host` — the same device->device in-place refresh EmbeddingKernel +// performs on its warm path, but with no host ids upload, so it can run +// INSIDE the captured region. `table_host` is the embedding table's host +// base (device shadow cached by the kEmbedding path). Requires WarmDecodeIds +// to have staged ids for this n first. +#ifdef VLLM_CPP_TENSTORRENT +void EmbedDeviceIdsInto(void* out_host, int64_t rows, int64_t cols, + const void* table_host, int64_t vocab, int64_t hidden, + int64_t n); +#else +inline void EmbedDeviceIdsInto(void*, int64_t, int64_t, const void*, int64_t, + int64_t, int64_t) {} +#endif + + + +#ifdef VLLM_CPP_TENSTORRENT +void WarmRopeCosSin(const int32_t* positions, int64_t tokens, int64_t hq, + int64_t hk, int64_t rot, double base); +#else +inline void WarmRopeCosSin(const int32_t*, int64_t, int64_t, int64_t, int64_t, + double) {} +#endif // ---- 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..ad6cd0353 100644 --- a/src/vt/tenstorrent/tenstorrent_ops.cpp +++ b/src/vt/tenstorrent/tenstorrent_ops.cpp @@ -21,6 +21,12 @@ #include "vt/cpu/cpu_threadpool.h" #include "vt/dtype.h" #include "vt/ops.h" +// This OBJECT library is not the `vllm` target, so it does not inherit the +// PUBLIC VLLM_CPP_TENSTORRENT define. Force the real declarations; the +// header's inline no-ops are only for CPU/Vulkan/Windows TUs. +#ifndef VLLM_CPP_TENSTORRENT +#define VLLM_CPP_TENSTORRENT +#endif #include "vt/tenstorrent/tenstorrent_device.h" #include @@ -51,6 +57,7 @@ #include #include #include +#include #include #include #include @@ -77,7 +84,15 @@ ttnn::Tensor chunked_scaled_dot_product_attention( #define VT_RESTORE_TRACY_ENABLE 1 #endif #include +#include #include +#include + +// Forward declare clone (header not in installed includes) +namespace ttnn { Tensor clone(const Tensor&, const std::optional&, const std::optional&, const std::optional&); } +#include +#include +#include // create_device_tensor, copy_to_device #include #include #include @@ -94,11 +109,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 +276,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 +310,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 +353,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 +512,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 +525,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 +554,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), @@ -466,17 +591,24 @@ bool TryDevicePagedUpdateBatch(ttnn::Tensor& cache_dev, MeshDevice& device, MakeHeightShardedUpdateInput(device, toks, base, C, nkv, nkv_pad, d, grid); std::vector pt(static_cast(C)); - std::vector update_idxs(static_cast(C)); + std::vector idxs(static_cast(C)); for (uint32_t b = 0; b < C; ++b) { pt[static_cast(b)] = static_cast(phys_blocks[static_cast(base + b)]); - update_idxs[static_cast(b)] = offsets[static_cast(base + b)]; + idxs[static_cast(b)] = static_cast(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); + // Paged mode requires update_idxs as a DEVICE tensor (the vector form + // alone is rejected: "Paged cache requires update_idxs tensor"). + ttnn::Tensor update_idxs_tensor = ttnn::Tensor::from_vector( + idxs, SpecOf(tt::tt_metal::Shape({static_cast(C)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); cache_dev = ttnn::experimental::paged_update_cache( - cache_dev, xt, update_idxs, /*update_idxs_tensor=*/std::nullopt, + cache_dev, xt, /*update_idxs=*/{}, update_idxs_tensor, /*share_cache=*/false, page_table, /*batch_offset=*/0, /*compute_kernel_config=*/std::nullopt, /*mesh_coords=*/std::nullopt); } @@ -491,43 +623,17 @@ bool TryDevicePagedFusedUpdateBatch(ttnn::Tensor& k_dev, ttnn::Tensor& v_dev, Me const std::vector& phys_blocks, const std::vector& offsets, const float* k_toks, const float* v_toks, uint32_t nkv, uint32_t d, - uint32_t /*bs*/) { + uint32_t /*bs*/) { const uint32_t B = static_cast(phys_blocks.size()); if (B == 0 || offsets.size() != phys_blocks.size()) return false; - try { - const uint32_t nkv_pad = std::max(32u, ((nkv + 31u) / 32u) * 32u); - const auto grid = device.compute_with_storage_grid_size(); - const uint32_t max_cores = - std::max(1u, static_cast(grid.x) * static_cast(grid.y)); - - for (uint32_t base = 0; base < B; base += max_cores) { - const uint32_t C = std::min(max_cores, B - base); - ttnn::Tensor xt_k = - MakeHeightShardedUpdateInput(device, k_toks, base, C, nkv, nkv_pad, d, grid); - ttnn::Tensor xt_v = - MakeHeightShardedUpdateInput(device, v_toks, base, C, nkv, nkv_pad, d, grid); - - std::vector pt(static_cast(C)); - std::vector update_idxs(static_cast(C)); - for (uint32_t b = 0; b < C; ++b) { - pt[static_cast(b)] = static_cast(phys_blocks[static_cast(base + b)]); - update_idxs[static_cast(b)] = offsets[static_cast(base + b)]; - } - ttnn::Tensor page_table = ttnn::Tensor::from_vector( - pt, SpecOf(tt::tt_metal::Shape({C, 1u}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), - &device); - - auto updated = ttnn::experimental::paged_fused_update_cache( - k_dev, xt_k, v_dev, xt_v, update_idxs, /*update_idxs_tensor=*/std::nullopt, - /*share_cache=*/false, page_table, /*batch_offset=*/0, - /*compute_kernel_config=*/std::nullopt, /*mesh_coords=*/std::nullopt); - k_dev = std::move(std::get<0>(updated)); - v_dev = std::move(std::get<1>(updated)); - } - return true; - } catch (const std::exception&) { - return false; - } + // The fused op rejects overlapping K/V sharded input grids + // ("input_tensor1 and input_tensor2 must not overlap"). Both K and V + // shards land on the same C cores via MakeHeightShardedUpdateInput. + // Fall back to two separate TryDevicePagedUpdateBatch calls (the paired + // path at TryDevicePagedPushPair handles this). + return false; + (void)k_dev; (void)v_dev; (void)device; (void)k_toks; (void)v_toks; + (void)nkv; (void)d; // suppress unused-param warnings } // Prefer fill for longer sequential prefills; otherwise batched update. @@ -662,6 +768,15 @@ void NotePagedKvRacWrites(Tensor& k_cache, Tensor& v_cache, const std::vector(cache_nhd.shape[3]); VT_CHECK(used_nb > 0 && used_nb <= pool_nb, "EnsurePagedKvTtnn: used_nb out of range"); - std::vector upload; { std::lock_guard g(PagedKvMutex()); PagedKvShadow& s = PagedKvShadows()[reinterpret_cast(cache_nhd.data)]; - if (s.device_current && s.device.has_value() && s.nb >= used_nb && s.nkv == nkv && s.bs == bs && s.d == d) { return *s.device; } + // Mirror grows to the full pool (zero-filled tail) so both mirror and + // device shadow stay at one stable size for the cache's lifetime. + EnsureMirrorCapacity(s, pool_nb, nkv, bs, d); + } - // Prefer incremental mirror; rebuild from NHD if cold/short/wrong geometry. + // Cold / short / geometry change: rebuild the used prefix from the host NHD + // cache into the (full-size) mirror; the tail stays zero. + std::vector used; + { + std::lock_guard g(PagedKvMutex()); + PagedKvShadow& s = PagedKvShadows()[reinterpret_cast(cache_nhd.data)]; if (!s.mirror_valid || s.nkv != nkv || s.bs != bs || s.d != d || s.nb < used_nb) { - // Release path: convert without holding the mutex for the whole NHD walk. - } else { - const size_t n_elems = static_cast(used_nb) * nkv * bs * d; - upload.assign(s.mirror.begin(), - s.mirror.begin() + static_cast(n_elems)); + s.mirror_valid = false; // content below [0,used) not trustworthy yet } } - - if (upload.empty()) { - // Cold / short mirror: full NHD→ttnn for the used prefix (outside the lock). - upload = NhdToTtnnLayoutPrefix(cache_nhd, used_nb); + used = NhdToTtnnLayoutPrefix(cache_nhd, used_nb); + { + std::lock_guard g(PagedKvMutex()); + PagedKvShadow& s = PagedKvShadows()[reinterpret_cast(cache_nhd.data)]; + EnsureMirrorCapacity(s, pool_nb, nkv, bs, d); + VT_CHECK(s.mirror.size() >= used.size(), "EnsurePagedKvTtnn: mirror shorter than used"); + std::memcpy(s.mirror.data(), used.data(), used.size() * sizeof(float)); + s.mirror_valid = true; } - const auto spec = SpecOf(tt::tt_metal::Shape({used_nb, nkv, bs, d}), ttnn::DataType::BFLOAT16, - ttnn::Layout::TILE); - ttnn::Tensor dev = ttnn::Tensor::from_vector(upload, spec, &device); + // Build the pool-sized device shadow WITHOUT a full-pool from_vector: the + // TILE-layout host transform is per-element and a 256-block vector costs + // seconds per cache (56 caches stalled cold for minutes). Upload only the + // used prefix, allocate the zero tail with ttnn::zeros (host std::fill + + // straight DMA — a constant fill is layout-order-agnostic), and stitch with + // a device-side concat. paged_update_cache / sdpa_decode never read the + // zero tail (page-table entries only cover allocated blocks). + const auto used_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"); + const auto dbg = std::getenv("VT_TT_TRACE_DEBUG") != nullptr; + const auto t0 = std::chrono::steady_clock::now(); + ttnn::Tensor dev = ttnn::Tensor::from_vector(used, used_spec, &device); + if (dbg) { + const auto t1 = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[TT-KVTIM] %p from_vector used_nb=%u %.0fms\n", + cache_nhd.data, used_nb, + std::chrono::duration(t1 - t0).count()); + } + if (used_nb < pool_nb) { + ttnn::Tensor tail = ttnn::zeros( + tt::tt_metal::Shape({pool_nb - used_nb, nkv, bs, d}), ttnn::DataType::BFLOAT16, + ttnn::Layout::TILE, std::ref(device)); + if (dbg) { + const auto t2 = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[TT-KVTIM] %p zeros tail_nb=%u %.0fms\n", + cache_nhd.data, pool_nb - used_nb, + std::chrono::duration(t2 - t0).count()); + } + dev = ttnn::concat(std::vector{dev, tail}, 0); + if (dbg) { + const auto t3 = std::chrono::steady_clock::now(); + std::fprintf(stderr, "[TT-KVTIM] %p concat pool_nb=%u %.0fms\n", + cache_nhd.data, pool_nb, + std::chrono::duration(t3 - t0).count()); + } + } std::lock_guard g(PagedKvMutex()); PagedKvShadow& s = PagedKvShadows()[reinterpret_cast(cache_nhd.data)]; - // If RAC raced and built a larger mirror, keep the larger one; still publish dev. - if (!s.mirror_valid || s.nkv != nkv || s.bs != bs || s.d != d || s.nb < used_nb) { - s.mirror = std::move(upload); - s.nb = used_nb; - s.nkv = nkv; - s.bs = bs; - s.d = d; - s.mirror_valid = true; - } s.device = dev; s.device_current = true; return dev; @@ -758,6 +906,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 +933,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 +962,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 +986,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 +1048,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 +1072,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 +1104,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); @@ -960,7 +1118,25 @@ void EmbeddingKernel(Queue&, Tensor& out, const Tensor& table, const Tensor& ids dev_out.logical_shape()[0] != t || dev_out.logical_shape()[1] != h) { dev_out = ttnn::reshape(dev_out, ttnn::Shape({t, h})); } - CommitDevice2D(out, std::move(dev_out)); + // HOST-FREE-DECODE: when the caller's buffer already carries a CURRENT + // device shadow of the same shape (the decode-graph driver's PERSISTENT + // hidden buffer), refresh that shadow IN PLACE (device->device copy) so + // its device address never moves. A captured region reads the address + // recorded at capture time; replacing the shadow here would leave every + // replay reading the capture-step embedding. First call (no shadow yet) + // commits normally. + bool in_place = false; + { + std::lock_guard g(SlotMutex()); + BufferSlot* s = FindSlot(out.data); + in_place = s != nullptr && s->device_current && s->device.has_value() && + s->dev_rows == t && s->dev_cols == h; + if (in_place) { + ttnn::copy(dev_out, *s->device); + s->host_current = false; + } + } + if (!in_place) CommitDevice2D(out, std::move(dev_out)); } // Upload a rank-1 affine vector as TILE BFLOAT16 [1, d], caching on the weight's @@ -981,6 +1157,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 +1182,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 +1223,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 +1248,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); @@ -1117,6 +1303,49 @@ void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& weight, CommitDevice2D(out, std::move(dev_y)); } +// kFusedChain: dispatch kFusedAddRmsNormStd to the same device RmsNorm path +// (residual += x; out = rms_norm(residual, weight)). Other recipes fall +// through to the CPU interpreter (host round-trip). Without this registration, +// FusedChain falls back to the CPU kernel which reads HOST memory — fatal +// when the PA output is device-resident (VT_TT_HOST_FREE_DECODE). +void FusedChainKernel(Queue& q, Tensor& out, const Tensor& x, const Tensor& weight, + Tensor* residual, const FusedRecipe& r, float eps) { + TT_OP_TRACE("FusedChain"); + // kFusedAddRmsNormStd: step0 kAdd(residual = x + residual), + // step1 kRmsNorm(out = rms_norm(residual, weight)). + // This is exactly RmsNormKernel with the residual parameter. + if (r.n == 2 && + r.steps[0].op == FOp::kAdd && r.steps[0].out == 2 && + r.steps[1].op == FOp::kRmsNorm && r.steps[1].out == 3 && + r.steps[1].gemma == false) { + RmsNormKernel(q, out, x, weight, RmsNormArgs{eps, false}, residual); + return; + } + // kFusedAddRmsNorm (gemma variant): same but gemma=true. + if (r.n == 2 && + r.steps[0].op == FOp::kAdd && r.steps[0].out == 2 && + r.steps[1].op == FOp::kRmsNorm && r.steps[1].out == 3 && + r.steps[1].gemma == true) { + RmsNormKernel(q, out, x, weight, RmsNormArgs{eps, true}, residual); + return; + } + // Unknown recipe: fall back to host (safe outside capture). + VT_CHECK(!tt_capture_active(), + "tenstorrent: unknown FusedChain recipe during capture"); + EnsureHost(out); + EnsureHost(x); + EnsureHost(weight); + if (residual != nullptr) EnsureHost(*residual); + // Delegate to the CPU interpreter by calling the registered CPU op. + auto cpu_fn = reinterpret_cast( + GetOpFallback(OpId::kFusedChain, DeviceType::kTENSTORRENT, "vt-tenstorrent")); + if (cpu_fn) { + cpu_fn(q, out, x, weight, residual, r, eps); + } else { + VT_CHECK(false, "tenstorrent: no FusedChain fallback available"); + } +} + // kSiluAndMul: SwiGLU gate half — out[i,j] = silu(x[i,j]) * x[i,j+d] // with d = x.shape[1]/2 (cpu_ops.cpp SiluAndMulKernel). Second Qwen3-dense // op beyond OPT (MLP: gate_up GEMM -> SiluAndMul -> down GEMM). Device path @@ -1124,6 +1353,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 +1386,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 +1399,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 +1468,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 +1610,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 +1635,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 +1737,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 +1804,337 @@ 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; + bool allocated = false; // ttnn::Tensor::is_allocated() crashes on default-constructed tensors in this build +}; +std::mutex& PaMetaMutex() { static std::mutex m; return m; } +std::map, PaMetaEntry>& PaMetaCache() { + static std::map, PaMetaEntry> c; + return c; +} +} // namespace + +// ---- R2: persistent cur_pos advanced on-device via plus_one ---------------- +// The PaMeta cur_pos tensor (read by sdpa_decode) and the RAC update_idxs +// tensor (read by paged_update_cache) both hold `seq_lens - 1` for decode. +// R2 aliases them: WarmDecodePos seeds the single persistent cur_pos tensor; +// CaptureDecodePosAdvance does plus_one on it inside the trace; WarmRacIdx's +// update_idxs copy_to_device is skipped (it reuses this tensor). +namespace { +struct DecodePosEntry { + ttnn::Tensor cur_pos; // int32 [num_reqs] device — advanced in-trace + bool allocated = false; +}; +std::mutex& DecodePosMutex() { static std::mutex m; return m; } +std::map& DecodePosCache() { + static std::map 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 (persistent, content refreshed) + ttnn::Tensor page_table; // int32 [C,1] device (persistent, content refreshed) + std::vector idx_host; // content identity for reuse check + // Persistent height-sharded RAC input: logical [1,1,nkv,d], padded + // [1,1,nkv_pad,d] (shard [nkv_pad,d] on one core). The in-region RAC + // ttnn::copy's the rope output into it; paged_update_cache reads only the + // first nkv rows (num_heads loop bound), so the padded tail rows are never + // read and may hold garbage — no zeros tail, no concat, no allocation. + ttnn::Tensor sharded_in; // K input (height-sharded) + ttnn::Tensor sharded_in_v; // V input (separate — K and V must NOT share the same buffer) + uint32_t nkv = 0; + uint32_t d = 0; + bool allocated = false; // ttnn::Tensor::is_allocated() crashes on default-constructed tensors in this build + bool sharded_in_is_alloc = false; +}; +std::mutex& RacIdxMutex() { static std::mutex m; return m; } +// Keyed by (num_slots, block_size) shape — idx tensors depend on slot values + block_size. +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; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && !tt_capture_active()) { + auto dump_src = [](const char* tag, const ttnn::Tensor& t) { + const auto ls = t.logical_shape(); + const auto ps = t.padded_shape(); + std::fprintf(stderr, + "[TT-TRACE] RAC src %s: logical=[", tag); + for (size_t i = 0; i < ls.rank(); ++i) + std::fprintf(stderr, "%s%u", i ? "," : "", ls[i]); + std::fprintf(stderr, "] padded=["); + for (size_t i = 0; i < ps.rank(); ++i) + std::fprintf(stderr, "%s%u", i ? "," : "", ps[i]); + std::fprintf(stderr, + "] dtype=%d layout=%d pages=%u strides0123=[%u,%u,%u,%u]\n", + (int)t.dtype(), (int)t.layout(), + t.buffer() ? t.buffer()->num_pages() : 0u, + ls.rank() > 0 ? t.strides()[0] : 0, ls.rank() > 1 ? t.strides()[1] : 0, + ls.rank() > 2 ? t.strides()[2] : 0, ls.rank() > 3 ? t.strides()[3] : 0); + }; + dump_src("k", *k_dev); + dump_src("v", *v_dev); + } + } + + // 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. + { + 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); + // WarmRacIdx (driver Refresh slot) refreshes update_idxs/page_table content + // every step via copy_to_device; here we just verify the tensors exist. + if (it == RacIdxCache().end() || !it->second.allocated) { + VT_CHECK(!tt_capture_active(), + "tenstorrent: RAC idx tensors not warmed — call WarmRacIdx " + "outside capture (driver Refresh slot) first"); + return false; + } + // sharded_in must exist (WarmRacIdx needs the paged-KV shadow geometry). + if (!it->second.sharded_in_is_alloc) { + VT_CHECK(!tt_capture_active(), + "tenstorrent: RAC sharded input not warmed — call WarmRacIdx " + "outside capture after WarmPagedKvShadow"); + return false; + } + } + + // Eager (cold) step AND capture: the identical op sequence (see + // build_input below). Running it eagerly first compiles the programs; + // capture then hits the program cache and replays against the persistent + // addresses. + RacIdxEntry rac_entry = [&] { + std::lock_guard g(RacIdxMutex()); + return RacIdxCache().at(std::make_pair(num_slots, static_cast(bs))); + }(); + + // Single code path for the eager (cold) step AND capture. The cold step must + // compile the exact programs the captured region will replay, so the op + // sequence and every TensorSpec must be identical in both phases: + // 1. reshape (metadata-only) to logical [1,1,nkv,d] + // 2. scalar multiply by 1.0 — eltwise ops allocate a FRESH output with a + // native 4D spec. Feeding the bare 2D→4D reshape view straight into + // ttnn::copy only writes head0 (the view's 2D-allocated storage + // confuses the tilized copy program), and a host-side to_vector → + // from_vector round-trip is illegal during capture and hashes + // differently (program-cache miss → binary load during capture). + // 3. ttnn::copy into the persistent sharded_in (preallocated output; the + // interleaved-TILE→height-sharded copy program uses only CBs) + // paged_fused_update_cache (in-place, has override_runtime_arguments) then + // ingests the sharded input against persistent addresses. + auto build_input = [&](const ttnn::Tensor& src, ttnn::Tensor& sharded_dst, + const RacIdxEntry& entry) -> ttnn::Tensor { + // Materialize a NATIVE [1,1,nkv,d] TILE tensor on device with a single + // code path on cold and capture (host round-trips are illegal during + // capture and compile differently-hashed programs). + // + // TILE readers map a logical element of a 4D [1,1,nkv,d] tensor to + // in-page ROW h of the (d/32)-page grid — so only [nkv,d]-SHAPED + // storage (heads at in-page rows 0..nkv-1) can be viewed; a [1,N] + // native (data at in-page row 0 of N/32 pages) mis-maps (head0-only + // or stale garbage — verified by per-head value dumps). Therefore: + // * [nkv,d] source (rope K output): explicit tile-padded 4D view — + // its spec is byte-identical to a native 4D's; multiply reads + // per-head exact (verified). + // * [1,N] source (QkvSplit V slice): materialize [nkv,d] storage + // first — per-head [1,d] slices concatenated on dim 0. Only ops + // already proven capture-safe in-region (slice/concat/eltwise). + // The scalar multiply materializes a fresh native 4D allocation; + // ttnn::copy moves it into the persistent sharded input. + const uint32_t nkv_pad = ((entry.nkv + 31u) / 32u) * 32u; + ttnn::Tensor laid_out = src; + if (src.logical_shape().rank() == 2 && src.logical_shape()[0] == 1) { + std::vector heads; + heads.reserve(entry.nkv); + for (uint32_t h = 0; h < entry.nkv; ++h) { + heads.push_back(ttnn::slice( + src, ttsl::SmallVector{0u, h * entry.d}, + ttsl::SmallVector{1u, (h + 1u) * entry.d}, + ttsl::SmallVector{1u, 1u})); + } + laid_out = ttnn::concat(heads, /*dim=*/0); + } + ttnn::Tensor native4 = ttnn::multiply( + ttnn::experimental::view( + laid_out, ttnn::Shape({1u, 1u, entry.nkv, entry.d}), + ttnn::Shape({1u, 1u, nkv_pad, entry.d})), + 1.0f); + auto head_maxima = [](const ttnn::Tensor& t, uint32_t nkv, uint32_t d) { + auto v = t.to_vector(); + std::string s; + for (uint32_t h = 0; h < nkv; ++h) { + float mx = 0; + for (uint32_t e = 0; e < d; ++e) { + const size_t i = static_cast(h) * d + e; + if (i < v.size()) mx = std::max(mx, std::abs(v[i])); + } + s += std::to_string(mx) + ","; + } + return s; + }; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr && !tt_capture_active()) { + auto chk = native4.to_vector(); + int nonzero = 0; + for (uint32_t h = 0; h < entry.nkv; ++h) { + float mx = 0; + for (uint32_t e = 0; e < entry.d; ++e) { + const size_t i = static_cast(h) * entry.d + e; + if (i < chk.size()) mx = std::max(mx, std::abs(chk[i])); + } + if (mx > 1e-6f) ++nonzero; + } + std::fprintf(stderr, "[TT-TRACE] RAC native4 nonzero_heads=%d/%u " + "src_headmax=[%s] laid_headmax=[%s] out_headmax=[%s]\n", + nonzero, entry.nkv, + head_maxima(src, entry.nkv, entry.d).c_str(), + head_maxima(laid_out, entry.nkv, entry.d).c_str(), + head_maxima(native4, entry.nkv, entry.d).c_str()); + } + ttnn::copy(native4, sharded_dst); + return sharded_dst; + }; + // V first, then K + // Debug: dump v_dev properties before sharding + ttnn::Tensor v_in = build_input(*v_dev, rac_entry.sharded_in_v, rac_entry); + ttnn::Tensor k_in = build_input(*k_dev, rac_entry.sharded_in, rac_entry); + // Debug: check v_in for all heads + // num_kv_heads_override pins the kernel's head loop to nkv rows: the input + // shard is tile-padded (nkv_pad rows) but only the first nkv rows hold data + // (upstream decode pattern, test_paged_cache_flexible_geometry.py). + // Use paged_fused_update_cache (single call for K+V) instead of two separate + // paged_update_cache calls. The fused op has override_runtime_arguments + // (the non-fused doesn't), so it works correctly with program cache enabled. + // The second separate call would reuse the first's cached program with the + // first's buffer addresses (program cache collision). + auto [new_kc, new_vc] = ttnn::experimental::paged_fused_update_cache( + *kc_dev, k_in, *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 device update (copy+paged_update_cache)\n"); + (void)offset; (void)block; + 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 +2220,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 +2274,44 @@ 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 { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA shadow miss: k_dc=%d k_dev=%d k_nb=%u/%u v_dc=%d v_dev=%d v_nb=%u/%u\n", + (int)sk.device_current, (int)sk.device.has_value(), sk.nb, used_nb, + (int)sv.device_current, (int)sv.device.has_value(), sv.nb, used_nb); + // 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 +2319,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) { @@ -1693,17 +2332,53 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso } ttnn::Tensor dev_q; bool q_from_device = false; - if (identity_q) { + if (identity_q) { try { - Tensor q_flat = query.View({total_q * hq, d}); - ttnn::Tensor dev_q2d = EnsureDevice2D(q_flat, device); - dev_q = ttnn::reshape(dev_q2d, ttnn::Shape({1u, Bu, hu, du})); + // SINGLE code path (cold compiles exactly what capture replays): + // resident [B*H, D] rope shadow → padding-correct 4D view → scalar + // multiply materializes a native [1, B, H, D] TILE tensor. The old + // Tensor::reshape view carried an UNPADDED spec (padded=logical) so + // sdpa_decode mis-mapped the storage (capture read head0-only while + // the cold host round-trip was correct); the explicit tile-padded + // view's spec is identical to a native 4D tensor's. + { + Tensor q_flat = query.View({total_q * hq, d}); + ttnn::Tensor dev_q_2d = EnsureDevice2D(q_flat, device); + const uint32_t hu_pad = ((hu + 31u) / 32u) * 32u; + dev_q = ttnn::multiply( + ttnn::experimental::view( + dev_q_2d, ttnn::Shape({1u, Bu, hu, du}), + ttnn::Shape({1u, Bu, hu_pad, du})), + 1.0f); + } + // Shard if needed + if (std::getenv("VT_TT_SHARD_Q") != nullptr) { + const uint32_t padded_hq = std::max(32u, hu); + const auto q_grid = device.compute_with_storage_grid_size(); + const tt::tt_metal::CoreRangeSet q_core_set = + tt::tt_metal::num_cores_to_corerangeset(Bu, q_grid, true); + tt::tt_metal::ShardSpec q_ss(q_core_set, {padded_hq, du}, + tt::tt_metal::ShardOrientation::ROW_MAJOR); + tt::tt_metal::MemoryConfig q_mc( + tt::tt_metal::TensorMemoryLayout::HEIGHT_SHARDED, + tt::tt_metal::BufferType::L1, q_ss); + dev_q = ttnn::to_memory_config(dev_q, q_mc); + } 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; } } if (!q_from_device) { + if (tt_capture_active()) { + VT_CHECK(false, "tenstorrent: PA Q host path is not capture-safe " + "(from_vector readback); the resident rope shadow " + "must be used during capture"); + } EnsureHost(query); std::vector q_host(static_cast(num_reqs * hq * d)); for (int64_t r = 0; r < num_reqs; ++r) { @@ -1715,32 +2390,185 @@ 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); - - const auto grid = device.compute_with_storage_grid_size(); - ttnn::operations::transformer::SDPAProgramConfig prog{ - grid, - std::nullopt, - /*q_chunk_size=*/32, - /*k_chunk_size=*/32, - /*exp_approx_mode=*/false, - /*max_cores_per_head_batch=*/16}; + ttnn::Tensor dev_pt, dev_pos; + bool use_warm_meta = false; + { + std::lock_guard g(PaMetaMutex()); + const auto pkey = std::make_pair(static_cast(num_reqs), + static_cast(max_blocks)); + auto it = PaMetaCache().find(pkey); + if (it != PaMetaCache().end() && it->second.allocated) { + dev_pt = it->second.page_table; + dev_pos = it->second.cur_pos; + const int32_t expect_cp = slens[0] - 1; // WarmPaMeta stores seq_lens - 1 + VT_CHECK(it->second.cp_host.size() >= 1 && it->second.cp_host[0] == expect_cp, + "tenstorrent: PA meta not warmed for this step"); + use_warm_meta = true; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA using cached meta (pt+cp) cap=%d\n", + (int)tt_capture_active()); + } + } + if (!use_warm_meta) { + // Trimmed inline page_table ([B, max_blocks_per_seq] like the + // upstream test): sdpa_decode mis-executes with a wide [B, 256] + // table when only a few blocks are allocated. + VT_CHECK(!tt_capture_active(), + "tenstorrent: PA meta not warmed for this step"); + const int64_t max_vblk = (num_reqs > 0 && slens[0] > 0) + ? (slens[0] - 1) / block_size + 1 : 1; + const int64_t pt_cols = std::min(max_blocks, std::max(2, max_vblk)); + std::vector pt_trim(static_cast(Bu * pt_cols)); + for (int64_t r = 0; r < num_reqs; ++r) { + for (int64_t c = 0; c < pt_cols; ++c) { + pt_trim[static_cast(r * pt_cols + c)] = + pt[static_cast(r * max_blocks + c)]; + } + } + dev_pt = ttnn::Tensor::from_vector( + pt_trim, SpecOf(tt::tt_metal::Shape({Bu, static_cast(pt_cols)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + 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); + } + // DON'T pass program_config — let sdpa_decode use its default. + // Our explicit config may interact badly with the program cache when + // called after other ops in the model forward. + + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] PA calling sdpa_decode cap=%d\n", (int)tt_capture_active()); + // Don't pass compute_kernel_config — let sdpa_decode use its default (HiFi2). + // Our HiFi4 was needed for nkv>1 correctness, but the 2-head issue is separate. + // Debug: when VT_TT_SDPA_TEST is set, create FRESH Q/KV/pt/pos from + // scratch (random data, all heads populated) and call sdpa_decode. + // This tests whether sdpa_decode works inside the model forward context + // with tensors created the same way as the Python standalone test. + if (std::getenv("VT_TT_SDPA_TEST") != nullptr && !tt_capture_active()) { + static bool tested = false; + if (!tested) { + tested = true; + fprintf(stderr, "[TT-SDPA-TEST] Running standalone sdpa_decode test inside model forward...\n"); + // Create fresh Q: [1,1,16,128] with random data in ALL heads + std::vector test_q(16 * 128); + for (auto& v : test_q) v = static_cast(rand()) / 2147483647.0f * 2.0f - 1.0f; + ttnn::Tensor test_q_rm = ttnn::Tensor::from_vector(test_q, + SpecOf(tt::tt_metal::Shape({1u, 1u, 16u, 128u}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), + nullptr); + ttnn::Tensor test_q_dev = test_q_rm.to_device(&device, tt::tt_metal::MemoryConfig{}); + ttnn::Tensor test_q_tile = ttnn::to_layout(test_q_dev, ttnn::Layout::TILE); + // Create fresh KV: [2,8,32,128] with data in block 1, offset 0 + std::vector test_k(2*8*32*128, 0.0f), test_v(2*8*32*128, 0.0f); + for (uint32_t h = 0; h < 8; h++) + for (uint32_t e = 0; e < 128; e++) { + size_t off = (1*8*32 + h*32 + 0) * 128 + e; + test_k[off] = static_cast(rand()) / 2147483647.0f * 2.0f - 1.0f; + test_v[off] = static_cast(rand()) / 2147483647.0f * 2.0f - 1.0f; + } + ttnn::Tensor test_k_rm = ttnn::Tensor::from_vector(test_k, + SpecOf(tt::tt_metal::Shape({2u,8u,32u,128u}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), nullptr); + ttnn::Tensor test_k_dev = test_k_rm.to_device(&device, tt::tt_metal::MemoryConfig{}); + ttnn::Tensor test_k_tile = ttnn::to_layout(test_k_dev, ttnn::Layout::TILE); + ttnn::Tensor test_v_rm = ttnn::Tensor::from_vector(test_v, + SpecOf(tt::tt_metal::Shape({2u,8u,32u,128u}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), nullptr); + ttnn::Tensor test_v_dev = test_v_rm.to_device(&device, tt::tt_metal::MemoryConfig{}); + ttnn::Tensor test_v_tile = ttnn::to_layout(test_v_dev, ttnn::Layout::TILE); + // page_table: [1,2] = [1,0] + std::vector test_pt = {1, 0}; + ttnn::Tensor test_pt_dev = ttnn::Tensor::from_vector(test_pt, + SpecOf(tt::tt_metal::Shape({1u,2u}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + // cur_pos: [1] = [0] + std::vector test_pos = {0}; + ttnn::Tensor test_pos_dev = ttnn::Tensor::from_vector(test_pos, + SpecOf(tt::tt_metal::Shape({1u}), ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + // Call sdpa_decode with FRESH Q + model KV/pt/pos + // But first check if dev_k has data at block 1 + { + auto k_check = dev_k.to_vector(); + const uint32_t k_nkv = dev_k.logical_shape()[1]; + const uint32_t k_bs = dev_k.logical_shape()[2]; + const uint32_t k_d = dev_k.logical_shape()[3]; + size_t b1_off = (1 * k_nkv * k_bs + 0 * k_bs + 0) * k_d; + fprintf(stderr, "[TT-SDPA-TEST] dev_k block1: [%f,%f,%f,%f] (off=%zu/%zu)\n", + k_check.size()>b1_off?k_check[b1_off]:0, + k_check.size()>b1_off+1?k_check[b1_off+1]:0, + k_check.size()>b1_off+2?k_check[b1_off+2]:0, + k_check.size()>b1_off+3?k_check[b1_off+3]:0, + b1_off, k_check.size()); + // Check V at block 1 for ALL 8 KV heads + auto v_check = dev_v.to_vector(); + const uint32_t v_nkv = dev_v.logical_shape()[1]; + const uint32_t v_bs = dev_v.logical_shape()[2]; + const uint32_t v_d = dev_v.logical_shape()[3]; + for (uint32_t h = 0; h < v_nkv; h++) { + size_t voff = (1 * v_nkv * v_bs + h * v_bs + 0) * v_d; + fprintf(stderr, "[TT-SDPA-TEST] dev_v block1 head%u: [%f,%f,%f,%f]\n", h, + v_check.size()>voff?v_check[voff]:0, + v_check.size()>voff+1?v_check[voff+1]:0, + v_check.size()>voff+2?v_check[voff+2]:0, + v_check.size()>voff+3?v_check[voff+3]:0); + } + // Also check if model KV WITHOUT RAC works: create fresh KV via + // to_layout(TILE) with the SAME data as dev_k + auto k_vec = dev_k.to_vector(); + auto v_vec = dev_v.to_vector(); + ttnn::Tensor fresh_k_rm = ttnn::Tensor::from_vector(k_vec, + SpecOf(tt::tt_metal::Shape({2u,8u,32u,128u}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), nullptr); + ttnn::Tensor fresh_k_dev = fresh_k_rm.to_device(&device, tt::tt_metal::MemoryConfig{}); + ttnn::Tensor fresh_k_tile = ttnn::to_layout(fresh_k_dev, ttnn::Layout::TILE); + ttnn::Tensor fresh_v_rm = ttnn::Tensor::from_vector(v_vec, + SpecOf(tt::tt_metal::Shape({2u,8u,32u,128u}), ttnn::DataType::BFLOAT16, ttnn::Layout::ROW_MAJOR), nullptr); + ttnn::Tensor fresh_v_dev = fresh_v_rm.to_device(&device, tt::tt_metal::MemoryConfig{}); + ttnn::Tensor fresh_v_tile = ttnn::to_layout(fresh_v_dev, ttnn::Layout::TILE); + // Call sdpa_decode with fresh Q + fresh-KV-from-model-data + ttnn::Tensor fresh_out = ttnn::transformer::paged_scaled_dot_product_attention_decode( + test_q_tile, fresh_k_tile, fresh_v_tile, dev_pt, + true, std::nullopt, dev_pos, std::nullopt, + 1.0f/std::sqrt(128.0f), std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt); + auto fresh_out_vec = fresh_out.to_vector(); + int fresh_nonzero = 0; + for (uint32_t h = 0; h < 16; h++) { + size_t off = static_cast(h) * 128; + float maxval = 0; + for (size_t i = off; i < off + 128 && i < fresh_out_vec.size(); i++) + maxval = std::max(maxval, std::abs(fresh_out_vec[i])); + if (maxval > 0.001f) fresh_nonzero++; + } + fprintf(stderr, "[TT-SDPA-TEST] fresh-KV-from-model-data: %d/16 heads\n", fresh_nonzero); + } + ttnn::Tensor test_out = ttnn::transformer::paged_scaled_dot_product_attention_decode( + test_q_tile, dev_k, dev_v, dev_pt, + true, std::nullopt, test_pos_dev, std::nullopt, + 1.0f/std::sqrt(128.0f), std::nullopt, std::nullopt, + std::nullopt, std::nullopt, std::nullopt, std::nullopt); + auto test_out_vec = test_out.to_vector(); + int non_zero = 0; + for (uint32_t h = 0; h < 16; h++) { + size_t off = static_cast(h) * 128; + float maxval = 0; + for (size_t i = off; i < off + 128 && i < test_out_vec.size(); i++) + maxval = std::max(maxval, std::abs(test_out_vec[i])); + if (maxval > 0.001f) non_zero++; + if (h % 2 == 0) + fprintf(stderr, "[TT-SDPA-TEST] head%u: max=%.4f %s\n", h, maxval, maxval > 0.001f ? "OK" : "ZERO"); + } + fprintf(stderr, "[TT-SDPA-TEST] Non-zero heads: %d/16\n", non_zero); + } + } + // Dump all shapes right before sdpa_decode ttnn::Tensor dev_out = ttnn::transformer::paged_scaled_dot_product_attention_decode( dev_q, dev_k, dev_v, dev_pt, /*is_causal=*/true, @@ -1749,12 +2577,16 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso /*attention_sink=*/std::nullopt, /*scale=*/args.scale, /*sliding_window_size=*/std::nullopt, - /*memory_config=*/std::nullopt, - /*program_config=*/prog, + /*memory_config=*/tt::tt_metal::MemoryConfig{}, + /*program_config=*/std::nullopt, /*compute_kernel_config=*/std::nullopt, /*paged_cache_geometry=*/std::nullopt, /*cache_position_modulo=*/std::nullopt); + // Dump PA output for comparison (first layer, cold step) + { + } + // Prefer keeping activations on device for o_proj: flatten to [B, H*D]. // Pure-decode with identity token order (qsl[r]==r) matches out's storage // layout [T,H,D] == [B,H,D] so Reshape→MatmulBT hits EnsureDevice2D. @@ -1767,9 +2599,13 @@ bool TryPagedAttentionDeviceDecode(Tensor& out, const Tensor& query, const Tenso } if (identity_order && total_q == num_reqs) { try { - ttnn::Tensor flat = - ttnn::reshape(dev_out, ttnn::Shape({Bu, static_cast(hq * d)})); - CommitDeviceLogical2D(out, std::move(flat), Bu, static_cast(hq * d)); + const uint32_t flat_cols = static_cast(hq * d); + ttnn::Tensor flat = ttnn::reshape(dev_out, + ttnn::Shape({Bu, flat_cols})); + CommitDeviceLogical2D(out, std::move(flat), Bu, flat_cols); + // Verify the committed output matches the PA output + { + } return true; } catch (const std::exception&) { // Fall through to host materialization. @@ -1778,6 +2614,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 +2630,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 +2748,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 +2769,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 +2798,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, @@ -2343,6 +3189,8 @@ struct Registrar { reinterpret_cast(static_cast(&PagedAttentionKernel))); RegisterOp(OpId::kGreedyArgmax, DeviceType::kTENSTORRENT, reinterpret_cast(static_cast(&GreedyArgmaxKernel))); + RegisterOp(OpId::kFusedChain, DeviceType::kTENSTORRENT, + reinterpret_cast(static_cast(&FusedChainKernel))); } } registrar; @@ -2364,6 +3212,37 @@ TraceState& TraceSlot() { return s; } constexpr auto kTraceCq = ttnn::QueueId(0); + +// Stall-bisection helper: counts completed TraceEndCaptureGraph calls. The +// bisection skip flags (VT_TT_NO_*_WARM) must NOT fire on the capture step +// itself — the captured rope cache-HIT guard requires fresh warm content — +// so they skip only once a graph exists (steady replay regime). +std::atomic& GraphCapturesCounter() { + static std::atomic n{0}; + return n; +} +int GraphCapturesDone() { return GraphCapturesCounter().load(); } +void NoteGraphCaptured() { GraphCapturesCounter().fetch_add(1); } +bool ReplayRegimeBisectSkip(const char* flag) { + return std::getenv(flag) != nullptr && GraphCapturesDone() > 0; +} + +// A replayed trace rewrote the device memory of every tensor the captured +// region produced, but the slot registry cannot know which host buffers those +// shadows belong to. Mark the host cache of EVERY device-current slot stale so +// the next host read re-downloads. Without this, DBuf::Download -> +// Backend::Copy -> EnsureHostBytes short-circuits on host_current and serves +// the bytes captured at trace time on every later replay (frozen logits). +// Replay is non-blocking, so device writes may still be in flight here; the +// invalidation only marks device memory as newer than the host copy, and the +// re-download at the next host read is the blocking sync point. +// Input shadows (weights, embeddings) are only re-read, never re-uploaded: +// replay does not modify them, and the extra download is identical bytes. +void InvalidateHostCachesAfterTrace() { + std::lock_guard g(SlotMutex()); + for (auto& [addr, slot] : Slots()) + if (slot.device_current && slot.device.has_value()) slot.host_current = false; +} } // namespace void TraceBeginCapture() { @@ -2372,6 +3251,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 +3271,7 @@ void TraceEndCapture() { s.replay_id = s.capturing_id; s.has_replay = true; s.capturing = false; + tt_capture_active() = false; } void TraceReplay() { @@ -2396,7 +3279,14 @@ void TraceReplay() { VT_CHECK(!s.capturing, "tenstorrent: TraceReplay during capture"); VT_CHECK(s.has_replay, "tenstorrent: TraceReplay with no captured trace"); MeshDevice& device = SharedMeshDevice(); - ttnn::operations::trace::execute_trace(&device, s.replay_id, kTraceCq, /*blocking=*/true); + // NON-BLOCKING, matching models/common/models/executor.py's long-decode + // pattern (execute_trace(blocking=False) + a later blocking readback): + // repeated blocking replays hang the mesh trace completion wait after a + // few dozen executions on this tt-metal build. The caller's post-replay + // device readback (logits Download) provides the synchronization; queue + // order keeps any later input refresh behind the replay. + ttnn::operations::trace::execute_trace(&device, s.replay_id, kTraceCq, /*blocking=*/false); + InvalidateHostCachesAfterTrace(); } void* TraceEndCaptureGraph() { @@ -2404,7 +3294,9 @@ void* TraceEndCaptureGraph() { VT_CHECK(s.capturing, "tenstorrent: TraceEndCaptureGraph without Begin"); MeshDevice& device = SharedMeshDevice(); ttnn::operations::trace::end_trace_capture(&device, s.capturing_id, kTraceCq); + NoteGraphCaptured(); s.capturing = false; + tt_capture_active() = false; // Opaque handle: heap-allocated MeshTraceId for the multi-graph API. return new ttnn::MeshTraceId(s.capturing_id); } @@ -2414,7 +3306,17 @@ void TraceReplayGraph(void* graph) { VT_CHECK(!TraceSlot().capturing, "tenstorrent: TraceReplayGraph during capture"); MeshDevice& device = SharedMeshDevice(); const auto id = *static_cast(graph); - ttnn::operations::trace::execute_trace(&device, id, kTraceCq, /*blocking=*/true); + // NON-BLOCKING — see TraceReplay: the qwen3 graph driver downloads the + // logits right after this call, and that blocking readback is the sync + // point (the upstream traced-decode executor's pattern). + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] execute_trace begin\n"); + ttnn::operations::trace::execute_trace(&device, id, kTraceCq, /*blocking=*/false); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] execute_trace enqueued\n"); + InvalidateHostCachesAfterTrace(); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] replay step complete\n"); } void TraceDestroyGraph(void* graph) { @@ -2423,11 +3325,129 @@ void TraceDestroyGraph(void* graph) { try { MeshDevice& device = SharedMeshDevice(); ttnn::operations::trace::release_trace(&device, *id); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] release_trace ok\n"); + } catch (const std::exception& ex) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] release_trace THREW: %s\n", ex.what()); } catch (...) { + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-STEP] release_trace THREW (unknown)\n"); } delete id; } +// ---- HOST-FREE-DECODE: persistent decode ids + capture-safe embedding ----- +// The replay step must perform ZERO eager device allocations: per-step eager +// alloc/free churn (the from_vector + embedding output of the old EmbedInto +// refresh) eventually hands a live trace's fixed buffer addresses to new +// allocations — tt-metal warns allocations while a trace exists "may be +// corrupted once a trace is executed", observed as a device hang ~60 replays +// in. The embedding therefore moves INSIDE the captured region: ids are +// refreshed into one persistent device tensor (allocation-free +// copy_to_device), the captured ttnn::embedding runs over that stable +// address, and its output tensor is kept alive so the trace's write address +// is never returned to the allocator. +namespace { +struct DecodeIdsEntry { + ttnn::Tensor ids; // device ROW_MAJOR UINT32 [n], content refreshed in place + ttnn::Tensor out; // embedding output [n, hidden] TILE; held for the trace + bool allocated = false; +}; +std::map& DecodeIdsCache() { + static std::map m; + return m; +} +std::mutex& DecodeIdsMutex() { + static std::mutex m; + return m; +} +} // namespace + +void WarmDecodeIds(const int32_t* ids, int64_t n) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + if (ids == nullptr || n < 1) return; + MeshDevice& device = SharedMeshDevice(); + std::vector host(static_cast(n)); + for (int64_t i = 0; i < n; ++i) { + VT_CHECK(ids[i] >= 0, "tenstorrent WarmDecodeIds: negative id"); + host[static_cast(i)] = static_cast(ids[i]); + } + const auto spec = SpecOf( + tt::tt_metal::Shape({static_cast(n)}), + ttnn::DataType::UINT32, ttnn::Layout::ROW_MAJOR); + std::lock_guard g(DecodeIdsMutex()); + DecodeIdsEntry& e = DecodeIdsCache()[n]; + const bool dbg_ids = std::getenv("VT_TT_TRACE_DEBUG") != nullptr; + if (dbg_ids) std::fprintf(stderr, "[TT-STEP] WarmDecodeIds begin n=%lld\n", (long long)n); + // VT_TT_NO_IDS_WARM: stall bisection only — skip the per-step H2D copy + // after the first capture (same token embedded every replay, numerically + // wrong, mechanics test only). + if (ReplayRegimeBisectSkip("VT_TT_NO_IDS_WARM")) { + if (dbg_ids) std::fprintf(stderr, "[TT-STEP] WarmDecodeIds skipped\n"); + return; + } + if (!e.allocated) { + e.ids = ttnn::Tensor::from_vector(host, spec, &device); + e.allocated = true; + } else { + // Allocation-free refresh: host staging tensor + H2D copy into the SAME + // device buffer (the WarmRacIdx pattern). + ttnn::Tensor h = ttnn::Tensor::from_vector(host, spec, nullptr); + ttnn::copy_to_device(h, e.ids); + } + if (dbg_ids) std::fprintf(stderr, "[TT-STEP] WarmDecodeIds done\n"); +} + +void EmbedDeviceIdsInto(void* out_host, int64_t rows, int64_t cols, + const void* table_host, int64_t vocab, int64_t hidden, + int64_t n) { + ttnn::Tensor dev_ids; + { + std::lock_guard g(DecodeIdsMutex()); + auto it = DecodeIdsCache().find(n); + VT_CHECK(it != DecodeIdsCache().end() && it->second.allocated, + "tenstorrent: EmbedDeviceIdsInto without WarmDecodeIds(n)"); + dev_ids = it->second.ids; + } + ttnn::Tensor dev_table; + { + std::lock_guard g(EmbedTableMutex()); + auto it = EmbedTableShadows().find(reinterpret_cast(table_host)); + VT_CHECK(it != EmbedTableShadows().end() && it->second.device.has_value() && + it->second.vocab == static_cast(vocab) && + it->second.h == static_cast(hidden), + "tenstorrent: EmbedDeviceIdsInto without a warmed embed table " + "(run one eager embedding step before capture)"); + dev_table = *it->second.device; + } + ttnn::Tensor dev_out = ttnn::embedding(dev_ids, dev_table, /*pad_token=*/std::nullopt, + /*layout=*/ttnn::Layout::TILE); + if (dev_out.logical_shape().rank() != 2 || + dev_out.logical_shape()[0] != n || dev_out.logical_shape()[1] != hidden) { + dev_out = ttnn::reshape( + dev_out, ttnn::Shape({static_cast(n), + static_cast(hidden)})); + } + { + std::lock_guard g(SlotMutex()); + BufferSlot* s = FindSlot(out_host); + VT_CHECK(s != nullptr && s->device_current && s->device.has_value() && + s->dev_rows == static_cast(rows) && + s->dev_cols == static_cast(cols), + "tenstorrent: EmbedDeviceIdsInto hidden shadow not resident"); + ttnn::copy(dev_out, *s->device); + s->host_current = false; + } + // Hold the embedding output for the trace's lifetime (its address is baked + // into the captured command sequence; freeing it would return the buffer + // to the allocator). + { + std::lock_guard g(DecodeIdsMutex()); + DecodeIdsCache()[n].out = dev_out; + } +} + // ---- Called from TenstorrentBackend::Alloc/Free/Copy (no ttnn in that TU). ---- void RegisterHostBuffer(void* host, size_t bytes) { if (host == nullptr) return; @@ -2467,6 +3487,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 +3522,541 @@ 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()) { + 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); + } else if (it->second.cos_host != ce) { + // In-place CONTENT refresh of the SAME device tensors: a captured rope + // op reads the address recorded at capture time, so replacing the + // tensors here would leave every replay reading the capture-step + // cos/sin (stale positions). The host tensors are built with the + // identical bf16 TILE spec so copy_to_device writes byte-matching + // data. Legal here: the driver calls this outside capture. + // VT_TT_NO_ROPE_REFRESH: stall bisection only — skip the per-step H2D + // copies AFTER the first capture (stale cos/sin on replays, numerically + // wrong, mechanics test only). + if (ReplayRegimeBisectSkip("VT_TT_NO_ROPE_REFRESH")) return; + ttnn::Tensor cos_h = ttnn::Tensor::from_vector( + ce, TileSpecOf(thu, halfu), nullptr); + ttnn::Tensor sin_h = ttnn::Tensor::from_vector( + se, TileSpecOf(thu, halfu), nullptr); + ttnn::copy_to_device(cos_h, it->second.cos); + ttnn::copy_to_device(sin_h, it->second.sin); + it->second.cos_host = ce; + } + }; + 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, + const int32_t* block_table, int64_t block_table_cols, + const int32_t* seq_lens) { + 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 sl0=%d\n", + (long long)num_slots, (long long)block_size, (long long)slots[0], + seq_lens ? seq_lens[0] : -1); + // VT_TT_NO_IDX_WARM: stall bisection only — skip the per-step H2D copies + // after the first capture (stale idx/page-table on device, numerically + // wrong, mechanics test only). + if (ReplayRegimeBisectSkip("VT_TT_NO_IDX_WARM")) return; + const bool r2_steady = std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr + && GraphCapturesDone() > 0; + if (num_slots < 1) return; + MeshDevice& device = SharedMeshDevice(); + // paged_update_cache needs: + // update_idxs[t] = the sequence position of the token being written + // (= seq_lens[t] - 1, the current decode position for user t) + // page_table = the block table (virtual→physical block mapping) + // The PA reads KV up to cur_pos = seq_lens - 1, so the RAC must write at + // exactly that position for the PA to see the current token's KV. + std::vector ptv; + std::vector idxv; + ptv.reserve(static_cast(num_slots)); + idxv.reserve(static_cast(num_slots)); + for (int64_t t = 0; t < num_slots; ++t) { + const int64_t slot = slots[t]; + if (slot < 0 || seq_lens == nullptr || block_table == nullptr) { + // Padding slot: paged_update_cache skips when update_idx == -1. + ptv.push_back(0); + idxv.push_back(-1); + continue; + } + // update_idx = the 0-indexed position of the token being decoded this step + // (= seq_lens[t] - 1, since seq_lens is the length BEFORE this token). + // paged_update_cache writes to page_table[vblk] * block_size + update_idx % block_size, + // which must equal the slot_mapping from the scheduler. + const int32_t cur_pos = seq_lens[t] - 1; + idxv.push_back(cur_pos); + // page_table = the physical block for cur_pos's virtual block. + const int32_t vblk = cur_pos / static_cast(block_size); + const int32_t pblk = block_table[t * block_table_cols + vblk]; + ptv.push_back(pblk); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmRacIdx user=%lld slot=%lld cur_pos=%d vblk=%d pblk=%d bt_cols=%lld (expect slot=%d)\n", + (long long)t, (long long)slot, cur_pos, vblk, pblk, (long long)block_table_cols, + pblk * static_cast(block_size) + cur_pos % static_cast(block_size)); + } + const auto key = std::make_pair(num_slots, block_size); + std::lock_guard g(RacIdxMutex()); + RacIdxEntry& e = RacIdxCache()[key]; + // idx/page-table tensors are allocated ONCE per key and their CONTENT is + // refreshed in place each step (copy_to_device, outside capture). The + // captured paged_update_cache replays against the stable address and reads + // the fresh values device-side. + ttnn::Tensor pt_host = ttnn::Tensor::from_vector( + ptv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots), 1u}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR)); + ttnn::Tensor idx_host = ttnn::Tensor::from_vector( + idxv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR)); + if (!e.allocated) { + e.page_table = ttnn::Tensor::from_vector( + ptv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots), 1u}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + // R2: alias update_idxs to the on-device-advanced cur_pos (DecodePosCache) + // when VT_TT_HOST_FREE_DECODE and the shapes match (decode T=1: + // num_slots == num_reqs). plus_one on cur_pos then advances update_idxs + // too, eliminating the per-replay update_idxs copy_to_device (toxic class). + bool aliased = false; + if (std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr) { + std::lock_guard dg(DecodePosMutex()); + auto dit = DecodePosCache().find(num_slots); + if (dit != DecodePosCache().end() && dit->second.allocated) { + e.update_idxs = dit->second.cur_pos; // share the same device buffer + aliased = true; + } + } + // After the first capture, a standalone update_idxs is never plus_one'd. + // Refuse rather than freeze the write index and emit fluent wrong tokens. + VT_CHECK(!r2_steady || aliased, + "tenstorrent: WarmRacIdx allocated a standalone update_idxs after " + "capture — plus_one will not advance it. Seed DecodePos per " + "cache entry; recapture does NOT clear this (#1105)."); + if (!aliased) { + e.update_idxs = ttnn::Tensor::from_vector( + idxv, SpecOf(tt::tt_metal::Shape({static_cast(num_slots)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), + &device); + } + e.allocated = true; + } else { + // R2 steady state: skip the per-step RAC page_table + update_idxs copies. + // update_idxs is aliased to the on-device-advanced cur_pos (plus_one). + // The RAC page_table (where paged_update_cache WRITES) only matters at + // block boundaries for decode T=1; for now skip in steady state (the PA + // page_table in WarmPaMeta handles the sdpa_decode read with on-change + // refresh). Phase 2 full: on-change refresh for RAC page_table too. + if (!r2_steady) { + ttnn::copy_to_device(pt_host, e.page_table); + ttnn::copy_to_device(idx_host, e.update_idxs); + } + } + e.idx_host = idxv; + // Build the persistent sharded RAC input ONCE from the first available + // paged-KV shadow's geometry (same nkv/d as the cache): logical + // [1,1,nkv,d], padded [1,1,nkv_pad,d], HEIGHT_SHARDED L1, shard + // [nkv_pad,d] on one core. paged_update_cache never reads the padded tail + // rows (num_heads loop bound), so it is left uninitialized — no zeros, no + // concat. + if (!e.sharded_in_is_alloc) { + 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 && (shadow.d % 32u) == 0u) { + const uint32_t np = std::max(32u, ((shadow.nkv + 31u) / 32u) * 32u); + 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); + // Logical [1,1,nkv,d]; TILE alignment derives the physical + // [1,1,nkv_pad,d] and the [nkv_pad,d] shard covers it on one core. + e.sharded_in = ttnn::create_device_tensor( + tt::tt_metal::TensorSpec( + tt::tt_metal::Shape({1u, 1u, shadow.nkv, shadow.d}), + tt::tt_metal::TensorLayout( + ttnn::DataType::BFLOAT16, + tt::tt_metal::PageConfig(ttnn::Layout::TILE), sm)), + &device); + // V needs a SEPARATE sharded buffer on a DIFFERENT core (forces a + // program cache miss so interleaved_to_sharded compiles a fresh + // program with V's buffer. Without different cores, the second call + // reuses K's cached program and writes to K's buffer). + const tt::tt_metal::CoreRangeSet core_set_v = + tt::tt_metal::CoreRangeSet({ + tt::tt_metal::CoreRange( + tt::tt_metal::CoreCoord(1, 0), + tt::tt_metal::CoreCoord(1, 0)) + }); + tt::tt_metal::ShardSpec ss_v(core_set_v, {np, shadow.d}, + tt::tt_metal::ShardOrientation::ROW_MAJOR); + tt::tt_metal::MemoryConfig sm_v( + tt::tt_metal::TensorMemoryLayout::HEIGHT_SHARDED, + tt::tt_metal::BufferType::L1, ss_v); + e.sharded_in_v = ttnn::create_device_tensor( + tt::tt_metal::TensorSpec( + tt::tt_metal::Shape({1u, 1u, shadow.nkv, shadow.d}), + tt::tt_metal::TensorLayout( + ttnn::DataType::BFLOAT16, + tt::tt_metal::PageConfig(ttnn::Layout::TILE), sm_v)), + &device); + e.nkv = shadow.nkv; + e.d = shadow.d; + e.sharded_in_is_alloc = true; + break; + } + } + } + // paged_update_cache + the interleaved-TILE→sharded ttnn::copy are 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 runs the identical copy+update sequence, compiling both programs. +} + +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)); + 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; + } + } + std::vector cpos(static_cast(num_reqs)); + for (int64_t r = 0; r < num_reqs; ++r) cpos[static_cast(r)] = seq_lens[r] - 1; + // Allocate ONCE per key; refresh CONTENT in place (copy_to_device, outside + // capture) so the captured sdpa_decode replays against a stable address + // while reading the fresh block-table/cur-pos values. + const auto key = std::make_pair(num_reqs, max_blocks); + // VT_TT_NO_IDX_WARM: legacy bisection override (skip ALL per-step copies). + if (ReplayRegimeBisectSkip("VT_TT_NO_IDX_WARM")) return; + const bool r2_steady = std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr + && GraphCapturesDone() > 0; + // R2 steady state: cur_pos/update_idxs advance on-device (plus_one); only + // page_table needs a host refresh, and only when it actually changed (block + // boundary crossed). The capture step (r2_steady==false) seeds everything. + ttnn::Tensor pt_host = 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)); + std::lock_guard g(PaMetaMutex()); + PaMetaEntry& e = PaMetaCache()[key]; + if (!e.allocated) { + e.page_table = 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); + // R2: alias cur_pos to the on-device-advanced DecodePos cur_pos (advanced + // by plus_one in the trace) when VT_TT_HOST_FREE_DECODE and it exists. + // sdpa_decode reads this tensor; plus_one advances it → no per-replay + // copy_to_device (the toxic ~38-replay hang class). + bool aliased = false; + if (std::getenv("VT_TT_HOST_FREE_DECODE") != nullptr) { + std::lock_guard dg(DecodePosMutex()); + auto dit = DecodePosCache().find(num_reqs); + if (dit != DecodePosCache().end() && dit->second.allocated) { + e.cur_pos = dit->second.cur_pos; // share the same device buffer + aliased = true; + } + } + // After the first capture, a standalone cur_pos is never plus_one'd. + // Refuse rather than freeze KV length and emit fluent wrong tokens. + VT_CHECK(!r2_steady || aliased, + "tenstorrent: WarmPaMeta allocated a standalone cur_pos after " + "capture — plus_one will not advance it. Seed DecodePos per " + "cache entry; recapture does NOT clear this (#1105)."); + if (!aliased) { + e.cur_pos = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + } + e.allocated = true; + } else { + // R2 steady state: page_table refreshes ONLY when content changed (block + // boundary crossed). cur_pos/update_idxs advance on-device via plus_one. + const bool pt_changed = (e.pt_host != pt); + if (pt_changed || !r2_steady) { + ttnn::copy_to_device(pt_host, e.page_table); + } + if (!r2_steady) { + ttnn::Tensor cp_host = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR)); + ttnn::copy_to_device(cp_host, e.cur_pos); + } + } + e.pt_host = pt; + e.cp_host = cpos; + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmPaMeta n=%lld mb=%lld cp0=%d r2=%d pt_chg=%d\n", + (long long)num_reqs, (long long)max_blocks, (int)cpos[0], + (int)r2_steady, (int)(e.pt_host != pt)); +} + +// R2: seed the persistent cur_pos device tensor (= seq_lens - 1) and warm the +// plus_one program (program cache) so CaptureDecodePosAdvance can run inside +// the trace. Called on the capture/warm step (re-seed), NOT every replay. +void WarmDecodePos(const int32_t* seq_lens, int64_t num_reqs) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + if (num_reqs < 1 || seq_lens == nullptr) return; + // R2: only seed/warm on the capture step (GraphCapturesDone()==0). On replay + // steps, cur_pos is advanced on-device by the captured plus_one — re-seeding + // here would overwrite the advance and break correctness. + // A new num_reqs after the first capture was never seeded: refuse rather + // than return and let WarmPaMeta/WarmRacIdx allocate a frozen standalone. + if (GraphCapturesDone() > 0) { + std::lock_guard g(DecodePosMutex()); + auto it = DecodePosCache().find(num_reqs); + VT_CHECK(it != DecodePosCache().end() && it->second.allocated, + "tenstorrent: WarmDecodePos after capture for a num_reqs that " + "was never seeded — cur_pos would freeze. Seed DecodePos per " + "cache entry; recapture does NOT clear this (#1105)."); + return; + } + MeshDevice& device = SharedMeshDevice(); + std::vector cpos(static_cast(num_reqs)); + for (int64_t r = 0; r < num_reqs; ++r) + cpos[static_cast(r)] = seq_lens[r] - 1; + + std::lock_guard g(DecodePosMutex()); + DecodePosEntry& e = DecodePosCache()[num_reqs]; + if (!e.allocated) { + e.cur_pos = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + e.allocated = true; + } else { + ttnn::Tensor cp_host = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR)); + ttnn::copy_to_device(cp_host, e.cur_pos); + } + // Warm plus_one (program cache) on a SCRATCH tensor so the in-trace call + // doesn't trigger "Cannot load new binaries during trace capture" — but + // leave e.cur_pos at its seeded value (the warm must NOT advance it, or the + // captured body reads cur_pos+1). + { + ttnn::Tensor scratch = ttnn::Tensor::from_vector( + cpos, SpecOf(tt::tt_metal::Shape({static_cast(num_reqs)}), + ttnn::DataType::INT32, ttnn::Layout::ROW_MAJOR), &device); + ttnn::operations::experimental::plus_one(scratch, + /*sub_core_grids=*/std::nullopt, /*skip_negative_entries=*/true); + } + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] WarmDecodePos n=%lld cp0=%d (seeded+warmed plus_one)\n", + (long long)num_reqs, (int)cpos[0]); +} + +// R2: capture ttnn::plus_one(cur_pos) at the END of the trace body. The NEXT +// replay sees cur_pos+1. Must be called INSIDE BeginCapture/EndCapture, after +// all reads of cur_pos (sdpa_decode / paged_update_cache) in the body. +void CaptureDecodePosAdvance(int64_t num_reqs) { + if (std::getenv("VT_TT_HOST_FREE_DECODE") == nullptr) return; + std::lock_guard g(DecodePosMutex()); + auto it = DecodePosCache().find(num_reqs); + if (it == DecodePosCache().end() || !it->second.allocated) { + std::fprintf(stderr, "[TT-TRACE] CaptureDecodePosAdvance: no seeded cur_pos for n=%lld\n", + (long long)num_reqs); + return; + } + ttnn::operations::experimental::plus_one(it->second.cur_pos, + /*sub_core_grids=*/std::nullopt, /*skip_negative_entries=*/true); + if (std::getenv("VT_TT_TRACE_DEBUG") != nullptr) + std::fprintf(stderr, "[TT-TRACE] CaptureDecodePosAdvance n=%lld (plus_one captured)\n", + (long long)num_reqs); +} + } // namespace vt::tenstorrent diff --git a/tests/scripts/test_agent_record.py b/tests/scripts/test_agent_record.py index 11385d176..a0df6bde6 100644 --- a/tests/scripts/test_agent_record.py +++ b/tests/scripts/test_agent_record.py @@ -1145,6 +1145,32 @@ def test_the_backend_pin_is_load_bearing_for_this_row(self) -> None: ) +class TenstorrentTraceRunnerRowIsCounted(TenstorrentMistralRowIsCounted): + """The BACKEND ratchet bump to 84 is backed by a real row (#1105).""" + + ROW = "BACKEND-TENSTORRENT-TRACE-RUNNER" + + def test_the_row_names_its_issue_and_its_spec(self) -> None: + text = (ROOT / ".agents/backend-matrix.md").read_text(encoding="utf-8") + row = next(l for l in text.splitlines() if l.startswith(f"| `{self.ROW}` |")) + self.assertIn("tenstorrent-trace-runner.md", row) + index = (ROOT / ".agents/issue-index.md").read_text(encoding="utf-8") + self.assertIn("issues/1105", index) + + +class TenstorrentHostFreeForwardRowIsCounted(TenstorrentMistralRowIsCounted): + """The BACKEND ratchet bump to 85 is backed by a real row (#1105).""" + + ROW = "BACKEND-TENSTORRENT-HOST-FREE-FORWARD" + + def test_the_row_names_its_issue_and_its_spec(self) -> None: + text = (ROOT / ".agents/backend-matrix.md").read_text(encoding="utf-8") + row = next(l for l in text.splitlines() if l.startswith(f"| `{self.ROW}` |")) + self.assertIn("tenstorrent-host-free-forward.md", row) + index = (ROOT / ".agents/issue-index.md").read_text(encoding="utf-8") + self.assertIn("issues/1105", index) + + class CudaLlamacppRowIsCounted(TenstorrentMistralRowIsCounted): """The BACKEND ratchet bump to 83 is backed by a real row (#979). 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); +}