feat(tenstorrent): Qwen3-dense ops, device residency, and paged-KV decode path - #240
Closed
lu-zero wants to merge 16 commits into
Closed
feat(tenstorrent): Qwen3-dense ops, device residency, and paged-KV decode path#240lu-zero wants to merge 16 commits into
lu-zero wants to merge 16 commits into
Conversation
…udler#1) FOLLOWING_AGENTS_PROTOCOL First Qwen3-dense (`Qwen3ForCausalLM`) op beyond OPT's set. Per-row RMS via ttnn::rms_norm (weight TILE [1,D], eps from RmsNormArgs); residual stream and gemma (w+1) stay host-staged to match cpu_ops RmsNormKernel. Default Qwen3 path (weight, no residual, gemma=false) is the device path. 1/1 new test case PASS on real Blackhole. Remaining: kSiluAndMul, Cast, kRopeNeox / kRopeCosSinCache / kRopeFromCache, then platform allow-list. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok-Code [Grok Build]
mudler#2) FOLLOWING_AGENTS_PROTOCOL SwiGLU gate half for Qwen3 MLP (gate_up -> SiluAndMul -> down). Host-staged bit-exact f32 path matching cpu_ops SiluAndMulKernel while Alloc is host memory; device ttnn::silu+mul deferred. 1/1 new test, 113/113 assertions PASS on Blackhole. Remaining: Cast, RoPE, platform allow-list. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok-Code [Grok Build]
FOLLOWING_AGENTS_PROTOCOL
Adds the remaining Qwen3-dense (`Qwen3ForCausalLM`) ops beyond OPT + RmsNorm
+ SiluAndMul, mirroring Metal M3b:
- kCastBf16 / kCastF32 (host Load/Store; K/V cache + logits paths)
- kRopeNeox (default dense RoPE; host f32 math, bit-exact)
- kRopeCosSinCache + kRopeFromCache (opt-in cache path; bit-identical to Neox)
Host-staged while Alloc is host memory. 3 new tests PASS on Blackhole
(Cast round-trip; RopeNeox bit-exact; CosSin+FromCache == Neox).
Next: TenstorrentPlatform supports_model_architecture("Qwen3ForCausalLM").
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Grok-Code [Grok Build]
FOLLOWING_AGENTS_PROTOCOL After landing RmsNorm, SiluAndMul, Cast, and RoPE (Metal M3b twin), the platform allow-list accepts Qwen3ForCausalLM alongside OPTForCausalLM. Other arches still fall back to CPU via SelectQueue. Platform test updated; full test_tenstorrent_backend green on Blackhole. Honest residual: e2e SACRED greedy on a small Qwen3 checkpoint not yet run in this commit. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok-Code [Grok Build]
FOLLOWING_AGENTS_PROTOCOL Keep ttnn results on-device between matmul/norm/add/relu ops via a host- pointer shadow map (Metal AllocMap shape): - RegisterHostBuffer/UnregisterHostBuffer on Alloc/Free - EnsureDevice2D uploads only when the device shadow is stale - CommitDevice2D publishes results without host download - EnsureHost / EnsureHostBytes (Backend::Copy) materialize for host ops and test reads; MarkHostWritten invalidates after host stores Host-staged ops (PA, RoPE, Silu, Cast, QkvSplit, …) still EnsureHost first. PagedAttention remains a host oracle. Unit tests: 16/16 PASS on Blackhole. Expected e2e win: skip host download+reupload on the dense matmul/norm path that dominated Qwen3 cost. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok-Code [Grok Build]
FOLLOWING_AGENTS_PROTOCOL - Extend the SACRED near-tie gate for kTENSTORRENT (device-aware our_ids / neartie_gap pair, op-provider proof, VT_DUMP_IDS bootstrap). - Commit Blackhole-captured our_ids_tenstorrent.npy and transformers teacher-forced neartie_gap_mnats_tenstorrent.npy (max gap 0.5 nats). - Add scripts/qwen3-neartie-gap-transformers.py for aarch64 hosts without a CUDA vLLM oracle venv. Full 16x16 gate re-run still host-PA bound; short smoke after residency: ~0.28 warm tok/s at 4 tokens. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Grok-Code [Grok Build]
Keep the gate_up → SiluAndMul → down GEMM chain device-resident: slice the last-dim halves, ttnn::silu(gate), multiply by up, CommitDevice2D. Relaxes the unit test from bit-exact host f32 to a BF16 envelope (same as norm/relu).
Host PA is still the long-context bottleneck (device sdpa_decode deferred). Optimize the oracle without changing the j-order / max-subtracted softmax: - hoist each query head into a local f32 buffer once - specialize f32/bf16/f16 loads to raw pointers (no per-elem dtype switch) - aarch64 NEON FMA for f32 dot and V axpy - parallelize over (token, head) only when nwork >= 64 (short decode stays serial — waking the 128-thread pool barriers dominated smoke) Short Qwen3-0.6B smoke stays ~0.28 warm tok/s (PA is tiny at seq~9; e2e is matmul/PCIe bound). PA microbench (T=1 Hq=16 D=128 seq=512) ~7 ms/call; enable with TT_PA_BENCH=1.
Implement NeoX rotary apply via ttnn slice + mul/sub/add + concat on a [T*H, D] view so a prior RmsNorm on that view can keep the device shadow. Short decode (T*H < 64) stays on a host apply path — many tiny device launches were slower than host at Qwen3-0.6B smoke shapes (~0.26 vs ~0.28 tok/s). Large prefill uses the device path. RopeCosSinCache remains host (tiny table). GPT-J interleave stays host-only. Unit tests (small T*H): bit-exact host path. Smoke warm ~0.28 tok/s flat.
When the merged QKV buffer already has a device shadow (typical post MatmulBT path), split with ttnn::slice on the last dim and CommitDevice2D each shard — avoids download+reupload before qk-RmsNorm. EnsureDevice2D now same-numel reshapes a resident shadow ([T, H*Dh] ↔ [T*H, Dh]) so the per-head RmsNorm view reuses QkvSplit's device result without a host round-trip. Host bit-exact memcpy remains for host-only qkv (unit tests). Device-path unit test: Relu→split within BF16 envelope. Short Qwen3 smoke warm still ~0.28 tok/s (other host barriers dominate).
Host PA still owns correctness; ttnn paged SDPA needs a different K/V layout ([nb,nkv,bs,d] vs our NHD [nb,bs,nkv,d]) so full device decode is deferred. Improvements on the host oracle: - Single-req causal path gathers short sequences into dense [seq,nkv,d] (≤32K floats) for sequential NEON + GQA reuse; long seq keeps paged walk (gather lost at seq=512: ~20ms vs ~7ms paged). - Prefill parallel uses a dedicated 4–16 thread pool instead of waking the 128-core global pool. Microbench seq=512: ~7.0 ms/call. Short Qwen3 smoke warm ~0.286–0.287 tok/s (small lift vs ~0.28).
Host KV stays vLLM NHD [nb,block,nkv,d] (LMCache/plane compatible). ReshapeAndCache still writes NHD and invalidates device shadows. On pure-decode PA (one query token per request, causal, TILE-legal D/block_size): - convert only physical blocks [0, max_page_id] to ttnn order [nb,nkv,bs,d] - upload TILE DRAM K/V shadows (reused until next RAC) - run ttnn::paged_scaled_dot_product_attention_decode - fall back to host oracle on shape/grid/exception Prefill and non-TILE geometry stay on the host path. Unit test (D=128,block=32,seq=64 decode) matches host within BF16 (~1e-4). Short smoke warm ~0.286 tok/s.
Keep a host-side ttnn-order mirror of each paged K/V cache. ReshapeAndCache patches only the written (block, offset) into the mirror and marks the device shadow dirty. EnsurePagedKvTtnn then re-uploads the used prefix from the mirror instead of re-walking the full NHD pool every layer. Cold/grow still rebuilds from NHD when the mirror is short or geometry mismatches. Non-TILE block/head dims drop shadows (host PA only). PA unit tests PASS; pure-decode BF16 path still ~1e-4 vs host. Short smoke warm ~0.288–0.289 tok/s.
… live After ReshapeAndCache patches the host ttnn-order mirror, try ttnn::experimental::paged_update_cache on the live K/V device shadows (B=1 height-sharded input, synthetic page_table mapping offset→phys block). Success keeps device_current so pure-decode PA reuses the shadow without re-upload. Failure falls back to mirror re-upload on the next Ensure. Workaround: undef TRACY_ENABLE around the experimental paged_cache include (metalium op_profiler expects 6-arg tracy_alloc_srcloc; tree has 5-arg). PA unit tests PASS; pure-decode BF16 ~1e-4. Short smoke warm ~0.28 tok/s.
Collect all valid ReshapeAndCache slots and push the live device shadow in one shot: paged_fill_cache for sequential prefill (T>=16 from offset 0), else height-sharded paged_update_cache with B tokens, chunked to the core grid. Host NHD + ttnn mirror stay authoritative; failures leave the shadow dirty so Ensure re-uploads from the mirror.
- Prefer paged_fused_update_cache for K+V RAC pushes (one launch); fall back to fill-pair then independent updates on failure. - Multi-token PA via chunked_scaled_dot_product_attention (per request, 32-token chunks, causal pad-safe). Decode stays on paged SDPA decode. - Unit test: pure-prefill T=32 matches host within BF16 envelope.
localai-bot
pushed a commit
that referenced
this pull request
Aug 10, 2026
feat(tenstorrent): Qwen3-dense ops, device residency, and paged-KV decode Follow-on to #197. Grows the Blackhole adapter from host-round-trip OPT ops into a Qwen3-dense path: kRmsNorm / kSiluAndMul / Cast + NeoX RoPE providers, device-resident shadows so the matmul chain skips the host download between ops, hybrid RoPE (device at T*H >= 64), and a dual-layout paged KV path for decode. Scope is tight — 9 files, all Tenstorrent-local except the Qwen3 gate and one analysis script — so the only real review question is whether it disturbs the devices that already work. It does not, on CUDA (GB10, sm_121a): test_qwen3_paged_engine 184/184 test_qwen27_paged_engine 235/235 (SACRED) That first one is the one that mattered. The gate's op-list assertion changes from kRopeCosSinCache + kRopeFromCache to kRopeNeox, which is an assertion about what the CUDA path runs, not just the Tenstorrent one. It passes, so the default CUDA path genuinely takes kRopeNeox with the cache off and the edit is a correction rather than a loosening. The Metal special-case generalises cleanly to : each partial accelerator is gated against ITS OWN oracle-backed golden with identical anchor+band logic, which is the right shape for a model whose 0.6B greedy stream is a genuine near-tie. Maintainer fix on top, authorship preserved: one commit carried no trailer block and two carried a malformed Assisted-by. Worth noting the missing block was REPORTED rather than crashing the checker, because 04069bd fixed that shadowed exact_waiver earlier today. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
Collaborator
|
Merged to main as Verified on CUDA (GB10, sm_121a), built from your branch rebased on main: That first gate is the one that mattered for review. Eight of the nine files are Tenstorrent-local, but The Maintainer fixes applied on top, authorship preserved as yours: one commit ( Nice piece of work: device-resident shadows so the matmul chain skips the host round-trip, the T*H>=64 hybrid RoPE threshold, and the dual-layout paged KV are exactly the seams that turn a proof-of-life adapter into a real decode path. |
Collaborator
|
Landed via local merge |
localai-bot
pushed a commit
that referenced
this pull request
Aug 10, 2026
Developer decision, on measured grounds rather than feel. Over the last 22 merged PRs the 900-line `product` budget was exceeded by 9 of them (41%): 71 76 78 129 178 182 237 338 340 479 596 605 895 | 1272 1306 1510 2261 2843 3224 4098 5282 6480 A gate that fires on four changes in ten is not a budget, it is noise that teaches people to waive it. Worse, tests are a third to a half of every large diff here (#211 909/2261, #197 684/1510, #240 931/2843, #196 2319/4098), so the budget charged RED-first mutation tests against the same allowance as kernel code, penalising exactly the discipline the rest of AGENTS.md demands. Size is now a review judgement: split a change when a reviewer would be better served by parts, not when a counter says so. WHAT IS NOT RETIRED, because dropping a size gate is not licence to drop the rules that shared its file: explicit path classification (no blanket directory exemptions), the fail-closed binary guard, the checker-change mutation-evidence contract, and the role check that keeps product paths on a PR. A new test pins all three so they cannot be deleted quietly alongside a constant. Evidence for this being a real checker-semantics change, per the contract this checker itself enforces: `test_no_line_budget_is_enforced_for_any_class` is RED against the pre-change checker (which exported PATH_CLASS_BUDGETS and rejected a 100k-line product change) and green after; the companion regression test is green on BOTH sides and asserts on the error, not its wording, so it is not coupled to a message the retirement reworded. Suite: 33 passed, 101 subtests. The `pr-size` CI job keeps its name because it is a required check; its comment and step name now say what it actually enforces. Re-running it over PR #211, the 2261-line change that prompted this, leaves exactly one error: the genuine missing mutation evidence for a checker edit. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude Code:claude-opus-5 [Claude Code]
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Follow-on to #197 (
BACKEND-TENSTORRENTOPT-125m W0–W2). Grows the Blackhole adapter from host-round-trip OPT ops into a Qwen3-dense path with selective device residency and a dual-layout paged KV path for decode.Qwen3-dense ops (host→hybrid device)
kRmsNorm,kSiluAndMul, Cast + NeoX RoPE providersQwen3ForCausalLMallowed onkTENSTORRENTResidency / hybrid compute
kSiluAndMul(slice + silu + mul)T·H ≥ 64, else host)kQkvSplitwhen qkv is already residentPaged attention / KV
[nb,nkv,bs,d]paged_scaled_dot_product_attention_decodewhen geometry allows (d%32==0,block%32==0)ReshapeAndCachepaged_update_cachewhen a shadow is live (single- and multi-token batch)paged_fill_cachefor longer sequential prefill (T≥16 from offset 0)Host NHD remains the source of truth for RAC / LMCache; device path is best-effort and falls back cleanly.
Test plan
tests/vt/test_tenstorrent_backend.cpp— 19/19 pass on real Blackhole (incl. device Silu/RoPE/QkvSplit envelopes, pure-decode PA path)--max-tokens 4 --repeat 3 --temperature 0viavllm-cli(warm ~0.28 tok/s; short run still matmul-bound)VLLM_CPP_TENSTORRENT)paged_fill_cachebatch path