feat(tenstorrent): map Backend graph capture onto ttnn mesh-trace - #354
Conversation
…dual/RoPE Keep the Blackhole hot path on device without re-introducing short-decode launch tax: - Persist embedding table (ROW_MAJOR BF16) and affine [1,D] weights; TILE embedding out via CommitDevice2D for the first residual/RMS/matmul. - Pure-decode PA commits [T, H·D] for o_proj; optionally reshapes a resident Q into [1,B,H,D]. - Residual RmsNorm: host for rows<32 (decode), device add+rms when larger; gemma stays host. RoPE: device only when T·H≥64. - Single-request multi-chunk prefill: permute+concat to [T,H·D] on device. Warm Qwen3-0.6B short smoke ~12.3 tok/s (was ~0.28 before embed cache). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok:grok-code [GrokBuild]
Wires SupportsGraphCapture/BeginCapture/EndCapture/Replay (and the multi-graph EndCaptureGraph/ReplayGraph/DestroyGraph handles) to ttnn::operations::trace begin/end/execute/release on the shared mesh device. Ops live in tenstorrent_ops.cpp so the backend TU stays free of ttnn headers. Contract matches CUDA graphs: warm the program cache first; keep device buffers fixed across capture/replay; no host Ensure/Download mid-capture. Unit test: warm matmul, capture, replay×3 (max_abs=0 vs warm), multi-graph handle path. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok:grok-code [GrokBuild]
444abb7 to
5956877
Compare
|
Landed on main, together with #289 which it stacks on — both as their own merge commits, your authorship intact. The split is what made these easy to take: the ttnn calls live in Landing scope, recorded in the merge commits so it is not overread later: the TU is its own OBJECT library behind Two author-declared deviations carried in verbatim from #289 rather than smoothed over: the device residual+RMS path at Reviewed and landed with Claude Code. |
BACKEND-TENSTORRENT-TRACE-RUNNER (child of BACKEND-TENSTORRENT). The graph-capture foundation landed as mudler#354 (SupportsGraphCapture()==true + BeginCapture/EndCapture/Replay + EndCaptureGraph/ReplayGraph/DestroyGraph -> ttnn trace, unit-tested). This spike decides whether wiring it into a capturable decode region is feasible and worth it. Read-only (POL-SPIKE-FIRST, POL-NO-CEILING); no code shipped. Three findings: Q1 (host-free region at T=1?): NO. Static trace: RmsNormKernel (rows=1<32) and PreferDeviceRope (T*H=16<64) both route to HOST every layer -- no contiguous device-only region spans a layer, so ttnn begin_trace_capture would abort on the host readback. Bonus: the shared decode-graph framework Qwen3DenseDecodeGraph (qwen3.cpp:489, used by Qwen3/Mistral/Llama/InternLM2) is DISABLED for TT because the platform does not override support_static_graph_mode() (base false; only CUDA returns true). Q2 (all-device-at-T=1 tok/s cost?): measured on real Blackhole P150, Qwen3-0.6B vllm-cli smoke, local-only VT_TT_FORCE_DEVICE override (reverted, unshipped): hybrid (current default): 12.5 tok/s warm all-device (forced): 10.7 tok/s warm Reproduces handoff §6 precisely (~12.3->10.5 there; 12.5->10.7 here; same ~14% / ~1.8 tok/s regression). So the decode-capture prerequisite costs ~1.8 tok/s; capture must beat 12.5 to be a net win -- implausible at T=1 where the host-API overhead capture collapses is small. Q3 (ttnn program-cache warm-up?): deferred -- moot until a host-free region exists and the perf case is made. DECISION: NO-GO for pure T=1 decode capture. The hybrid thresholds are load-bearing for perf (Q2), and the platform support_static_graph_mode() flag must NOT flip to true until a capture scope that beats eager exists (flipping it enables the decode-graph framework, which would try to capture the T=1 path and abort or regress). Open follow-on (separate row): prefill / multi-token chunk capture, where thresholds already go all-device -- that region IS host-free and is the natural first capture target. Surveys how the other backends structure capture: CUDA is the ONLY SupportsGraphCapture()==true backend (cuda_backend.cu:184-240); Metal/Vulkan are false with the eventual mapping named but unimplemented. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Maki:glm-5.2 [Maki]
2 participants
Summary
Draft. Stacked on #289 (residency chain). Do not merge before #289 lands; rebase onto
mainafter.Maps
vt::Backendgraph capture onto ttnn mesh-trace so decode can amortize host launch tax the same way CUDA graphs do.API surface
SupportsGraphCapture() == trueBeginCapture/EndCapture/Replay→ single-slot mesh traceEndCaptureGraph/ReplayGraph/DestroyGraph→ opaque multi-graph handlesImplementation lives in
tenstorrent_ops.cpp(freeTrace*functions) so the backend TU stays free of ttnn headers.Contract (same class as CUDA graphs)
BeginCaptureEnsure/Download/ fresh program compiles mid-captureThis PR is one commit on top of #289's tip. GitHub cannot set the base to the fork's residency branch against
mudler/vllm.cpp, so the PR still targetsmainbut the history is rebased on residency (0365259b+ this commit).Test plan
max_abs=0vs warm, multi-graph handle)