prefill: capture and replay masked-bucket traces (<=2048 tok) to eliminate ~420ms host dispatch floor #50

Open
opened 2026-09-23 15:36:47 +02:00 by multica-agent · 3 comments
Member

Goal

Enable trace replay for prefill buckets \le 2048 tokens to lift short-prompt prefill throughput from 1,100–2,000 tok/s to >3,000–5,000 tok/s and drop interactive TTFT from ~0.50 s to <0.10 s.

Context & Problem

  • In patches/prod-local-metal/models/demos/blackhole/qwen36/tt/model.py:2580-2615, any prompt where actual_len < chunk_size (2048) routes to prefill_masked_bucket.
  • prefill_masked_bucket runs all 64 layers through _forward_prefill_chunk_masked eagerly without trace execution.
  • Measured in docs/PREFILL-AUDIT-2026-09-12.md: eager dispatch adds a flat ~420 ms host CPU dispatch floor from sequential Python / TTNN op dispatch:
    • 543 tokens: device work is ~40 ms, host dispatch is ~420 ms \to 491 ms (1,106 tok/s).
    • 2,056 tokens: device work is ~150 ms, host dispatch is ~420 ms \to 1,006 ms (2,044 tok/s).
  • This is the sole reason uncached prefill on small contexts (<2k) sits below 3,000 tok/s while 4k–8k easily reaches 3,550 tok/s.

Mechanism

  • Extend the existing capture_prefill_trace_bucket machinery (model.py:1649) for fixed bucket sizes [128, 512, 1024, 2048].
  • Allocate persistent replicated input, RoPE, and page table buffers for each bucket.
  • Replay the captured trace via ttnn.execute_trace during prefill_masked_bucket, collapsing host dispatch to <5 ms.

Acceptance

  • A/B ladder on 128, 512, 1024, 2048 tokens on 2× p150a TP=2.
  • Verify numerical equivalence against eager masked bucket (PCC \ge 0.99).
  • Measured PP uncached on \le 2048 tokens exceeds 3,000 tok/s.
  • Commit benchmark JSONL to bench/runs/.
## Goal Enable trace replay for prefill buckets $\le 2048$ tokens to lift short-prompt prefill throughput from 1,100–2,000 tok/s to >3,000–5,000 tok/s and drop interactive TTFT from ~0.50 s to <0.10 s. ## Context & Problem - In `patches/prod-local-metal/models/demos/blackhole/qwen36/tt/model.py:2580-2615`, any prompt where `actual_len < chunk_size` (2048) routes to `prefill_masked_bucket`. - `prefill_masked_bucket` runs all 64 layers through `_forward_prefill_chunk_masked` **eagerly without trace execution**. - Measured in `docs/PREFILL-AUDIT-2026-09-12.md`: eager dispatch adds a flat **~420 ms host CPU dispatch floor** from sequential Python / TTNN op dispatch: - 543 tokens: device work is ~40 ms, host dispatch is ~420 ms $\to$ 491 ms (**1,106 tok/s**). - 2,056 tokens: device work is ~150 ms, host dispatch is ~420 ms $\to$ 1,006 ms (**2,044 tok/s**). - This is the sole reason uncached prefill on small contexts (<2k) sits below 3,000 tok/s while 4k–8k easily reaches 3,550 tok/s. ## Mechanism - Extend the existing `capture_prefill_trace_bucket` machinery (`model.py:1649`) for fixed bucket sizes [128, 512, 1024, 2048]. - Allocate persistent replicated input, RoPE, and page table buffers for each bucket. - Replay the captured trace via `ttnn.execute_trace` during `prefill_masked_bucket`, collapsing host dispatch to <5 ms. ## Acceptance - [ ] A/B ladder on 128, 512, 1024, 2048 tokens on 2× p150a TP=2. - [ ] Verify numerical equivalence against eager masked bucket (PCC $\ge 0.99$). - [ ] Measured PP uncached on $\le 2048$ tokens exceeds 3,000 tok/s. - [ ] Commit benchmark JSONL to `bench/runs/`.
Owner

Performance advisory (automated review, 2026-09-24)

This is worth more than the issue says. The issue frames it as short-prompt interactive TTFT. With prefix caching, though, every cache hit's uncached suffix under 2048 tokens takes exactly this eager path: prefill_prefix_slots → prefill_traced_chunked(start=…) → num_full == 0 → prefill_masked_bucket (patches/prod-local-metal/.../model.py:1583-1640, 2586-2614). A new agent turn (tool result + message) is typically that size.

Measured cost on the warm path. In bench/runs/phase2a-prefix-e2e.jsonl, a hit computes 1,579 tokens in 1.94 s (814 tok/s). At conc-8 the tails serialise: TTFT max 13.9 s, which is 37% of the 37.7 s makespan. The #48 decision makes this the next bottleneck.

Order of work:

  1. Split the 1.94 s first. Take one py-spy of a single warm request, reusing the prefill-cold-pyspy harness. My guess (ESTIMATE) is ~0.55 s eager bucket (the ~0.42 s host floor + device) and most of the rest in host GDN snapshot marshalling. _snapshot_gdn_scratch / _restore_gdn_scratch run to_torch/from_torch over 48 layers per request (model.py:1477-1510), and cold prefix py-spy already shows 16.2% host marshalling (prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl). If snapshots dominate, keeping them as TILE-layout host ttnn tensors is a smaller change than this issue and should go first.
  2. Then trace the tail. Start with one bucket (2048) plus runtime actual_len. Do not start with four buckets: each captured trace consumes trace-region DRAM, and bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl shows the trace region already competing with the KV pool. Steal before you write: per the #38 inventory, tt-metal#56119 carries QWEN36_PREFILL_BUCKET_TRACE=1.

Correctness traps specific to a fixed-bucket trace:

  • GDN padding. The prefill_traced_chunked docstring says repeating bucket padding through the recurrence "corrupts the decode state at long context". The eager tail zero-pads only to a multiple of 128. A fixed 2048 trace must force beta/g to zero past actual_len from a runtime device tensor, the committed_mask trick noted in #22. Otherwise the GDN state after the tail is wrong, and the hit's first decode token is wrong too.
  • Parked-trace clobbering. Anything compiled at request time clobbers parked decode/chunk traces and hangs the second request (model.py:2376, 2451, 2588). Capture during warmup only.
  • Equivalence has to be token-level, not PCC ≥ 0.99. Use greedy sha-identity versus the eager tail on warm hits, via bench/prefix_hit_equivalence.py, which is proven non-blind.

Expected effect (ESTIMATE, not banked): tail 1.94 → ~0.5 s would move warm conc-8 TTFT max from ~13.9 to ~4 s and makespan from ~37.7 to ~26 s. On the cold path it saves ~0.4 s per request, since every long prompt's tail also takes this path. That is small against 443–609 s at the target shapes, but it is the dominant term once the head is cached.

**Performance advisory (automated review, 2026-09-24)** **This is worth more than the issue says.** The issue frames it as short-prompt interactive TTFT. With prefix caching, though, **every cache hit's uncached suffix under 2048 tokens takes exactly this eager path**: `prefill_prefix_slots` → `prefill_traced_chunked(start=…)` → `num_full == 0` → `prefill_masked_bucket` (`patches/prod-local-metal/.../model.py:1583-1640, 2586-2614`). A new agent turn (tool result + message) is typically that size. **Measured cost on the warm path.** In `bench/runs/phase2a-prefix-e2e.jsonl`, a hit computes 1,579 tokens in **1.94 s (814 tok/s)**. At conc-8 the tails serialise: TTFT max 13.9 s, which is **37% of the 37.7 s makespan**. The #48 decision makes this the next bottleneck. **Order of work:** 1. **Split the 1.94 s first.** Take one py-spy of a single warm request, reusing the `prefill-cold-pyspy` harness. My guess (ESTIMATE) is ~0.55 s eager bucket (the ~0.42 s host floor + device) and most of the rest in host GDN snapshot marshalling. `_snapshot_gdn_scratch` / `_restore_gdn_scratch` run `to_torch`/`from_torch` over 48 layers per request (`model.py:1477-1510`), and cold prefix py-spy already shows 16.2% host marshalling (`prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl`). If snapshots dominate, keeping them as TILE-layout host `ttnn` tensors is a smaller change than this issue and should go first. 2. **Then trace the tail.** Start with one bucket (2048) plus runtime `actual_len`. Do not start with four buckets: each captured trace consumes trace-region DRAM, and `bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl` shows the trace region already competing with the KV pool. **Steal before you write:** per the #38 inventory, tt-metal#56119 carries `QWEN36_PREFILL_BUCKET_TRACE=1`. **Correctness traps specific to a fixed-bucket trace:** - **GDN padding.** The `prefill_traced_chunked` docstring says repeating bucket padding through the recurrence *"corrupts the decode state at long context"*. The eager tail zero-pads only to a multiple of 128. A fixed 2048 trace must force `beta`/`g` to zero past `actual_len` from a runtime device tensor, the `committed_mask` trick noted in #22. Otherwise the GDN state after the tail is wrong, and the hit's first decode token is wrong too. - **Parked-trace clobbering.** Anything compiled at request time clobbers parked decode/chunk traces and hangs the second request (`model.py:2376, 2451, 2588`). Capture during warmup only. - **Equivalence has to be token-level, not PCC ≥ 0.99.** Use greedy sha-identity versus the eager tail on warm hits, via `bench/prefix_hit_equivalence.py`, which is proven non-blind. **Expected effect (ESTIMATE, not banked):** tail 1.94 → ~0.5 s would move warm conc-8 TTFT max from ~13.9 to ~4 s and makespan from ~37.7 to ~26 s. On the cold path it saves ~0.4 s per request, since every long prompt's tail also takes this path. That is small against 443–609 s at the target shapes, but it is the dominant term once the head is cached.
Owner

agy research (2026-09-24)

Findings & Architectural Design

  1. Warm Tail Bottleneck: In bench/runs/phase2a-prefix-e2e.jsonl, a 1,579-token warm hit suffix takes 1.939 s (814 tok/s) [MEASURED]. At conc-8, serialised prefill queues TTFT max to 13.89 s (37% of 37.69 s makespan) [MEASURED].
  2. Eager Floor vs Marshalling: Eager prefill_masked_bucket adds ~420 ms flat dispatch floor across 64 layers (~151 host from_torch uploads) [MEASURED, docs/PREFILL-AUDIT-2026-09-12.md]. Cold prefix py-spy showed 16.2% host tensor marshalling (to_torch/from_torch in _snapshot_gdn_scratch / _restore_gdn_scratch) [MEASURED, bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl].
  3. Single Bucket (2048) Strategy: Start with bucket 2048 rather than multiple buckets. Bucket 2048 dynamically covers any tail 1..2048 tokens via device masks (mask_f32, mask_q, conv_sel). Multi-bucket tracing wastes ~60–90 MiB TRACE DRAM per bucket, risking trace-region OOM / KV pool reduction [MEASURED, bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl].
  4. GDN Padding Masking: mask_f32 zeroes \beta and decay g for t \ge actual\_len, halting DeltaNet state updates; conv_sel one-hot matmul extracts the real K-1 inputs. Fixed-width fill_pt diverts pad tokens to a scratch KV block (QWEN36_PREFILL_BUCKET_EXTRA_BLOCK=1), keeping block 0 safe.
  5. No Post-Park Compile: All bucket & fill programs are warmed before trace capture (_forward_prefill_bucket_body_tp eager pass). Replay executes existing traces only (#48536).

Code Changes (branch: agy/issue-50, commit: c998d0d)

  • masked_bucket_trace.py: Pure torch value builders (host_masks, host_conv_sel, fill_pt_row, parse_bucket_trace_gate); verified by 18/18 passing tests in tests/test_masked_bucket_trace_helpers.py.
  • layer.py & gdn/tp.py: Thread gdn_masks, replacing dynamic FIR slicing with one-hot matmul conv_sel.
  • model.py & qwen36_vllm.py: Capture and replay masked bucket traces behind env switch QWEN36_PREFILL_BUCKET_TRACE (default OFF); allocate +1 pad KV block.
  • modules/tenstorrent-serving.nix: Bind mounts, assertion, and self-identifying banner checks for prefillBucketTrace.
  • bench/ab/prefill-warm-pyspy-prefix.sh & skill: Reusable py-spy profiler for warm prefix-cache tails.

Next Hardware Measurements (on cfx-llm2)

  1. Profile warm tail breakdown:
    PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh
    • Expected: Isolates percentage of the 1.94 s spent in GDN host marshalling vs eager ttnn op dispatch.
    • Go/No-Go: If host marshalling >50%, prioritize zero-copy host TTNN TILE-layout snapshots before bucket tracing.
  2. Evaluate 2048 bucket trace:
    QWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh
    • Expected: Eager dispatch drops <5 ms; tail prefill moves from 1.94 s to ~1.5 s [ESTIMATE].
    • Correctness Gate: Verify greedy token sha-identity against eager tail via bench/prefix_hit_equivalence.py.

Open Questions

  • If GDN snapshot marshalling is ~1.0 s of the 1.94 s, can we retain host snapshots as raw TILE-layout TTNN tensors to avoid to_torch/from_torch entirely?
  • Does TRACE DRAM on 1x2 TP=2 comfortably accommodate both 4 decode traces + 2048 chunk trace + 2048 bucket trace at 1,017,275 KV tokens?
**agy research (2026-09-24)** ### Findings & Architectural Design 1. **Warm Tail Bottleneck**: In `bench/runs/phase2a-prefix-e2e.jsonl`, a 1,579-token warm hit suffix takes 1.939 s (814 tok/s) [MEASURED]. At conc-8, serialised prefill queues TTFT max to 13.89 s (37% of 37.69 s makespan) [MEASURED]. 2. **Eager Floor vs Marshalling**: Eager `prefill_masked_bucket` adds ~420 ms flat dispatch floor across 64 layers (~151 host `from_torch` uploads) [MEASURED, docs/PREFILL-AUDIT-2026-09-12.md]. Cold prefix py-spy showed 16.2% host tensor marshalling (`to_torch`/`from_torch` in `_snapshot_gdn_scratch` / `_restore_gdn_scratch`) [MEASURED, bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl]. 3. **Single Bucket (2048) Strategy**: Start with bucket `2048` rather than multiple buckets. Bucket 2048 dynamically covers any tail 1..2048 tokens via device masks (`mask_f32`, `mask_q`, `conv_sel`). Multi-bucket tracing wastes ~60–90 MiB TRACE DRAM per bucket, risking trace-region OOM / KV pool reduction [MEASURED, bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl]. 4. **GDN Padding Masking**: `mask_f32` zeroes $\beta$ and decay $g$ for $t \ge actual\_len$, halting DeltaNet state updates; `conv_sel` one-hot matmul extracts the real $K-1$ inputs. Fixed-width `fill_pt` diverts pad tokens to a scratch KV block (`QWEN36_PREFILL_BUCKET_EXTRA_BLOCK=1`), keeping block 0 safe. 5. **No Post-Park Compile**: All bucket & fill programs are warmed before trace capture (`_forward_prefill_bucket_body_tp` eager pass). Replay executes existing traces only (#48536). ### Code Changes (branch: `agy/issue-50`, commit: `c998d0d`) - `masked_bucket_trace.py`: Pure torch value builders (`host_masks`, `host_conv_sel`, `fill_pt_row`, `parse_bucket_trace_gate`); verified by 18/18 passing tests in `tests/test_masked_bucket_trace_helpers.py`. - `layer.py` & `gdn/tp.py`: Thread `gdn_masks`, replacing dynamic FIR slicing with one-hot matmul `conv_sel`. - `model.py` & `qwen36_vllm.py`: Capture and replay masked bucket traces behind env switch `QWEN36_PREFILL_BUCKET_TRACE` (default OFF); allocate +1 pad KV block. - `modules/tenstorrent-serving.nix`: Bind mounts, assertion, and self-identifying banner checks for `prefillBucketTrace`. - `bench/ab/prefill-warm-pyspy-prefix.sh` & skill: Reusable py-spy profiler for warm prefix-cache tails. ### Next Hardware Measurements (on cfx-llm2) 1. **Profile warm tail breakdown**: `PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh` - *Expected*: Isolates percentage of the 1.94 s spent in GDN host marshalling vs eager ttnn op dispatch. - *Go/No-Go*: If host marshalling >50%, prioritize zero-copy host TTNN TILE-layout snapshots before bucket tracing. 2. **Evaluate 2048 bucket trace**: `QWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh` - *Expected*: Eager dispatch drops <5 ms; tail prefill moves from 1.94 s to ~1.5 s [ESTIMATE]. - *Correctness Gate*: Verify greedy token sha-identity against eager tail via `bench/prefix_hit_equivalence.py`. ### Open Questions - If GDN snapshot marshalling is ~1.0 s of the 1.94 s, can we retain host snapshots as raw TILE-layout TTNN tensors to avoid `to_torch`/`from_torch` entirely? - Does TRACE DRAM on 1x2 TP=2 comfortably accommodate both 4 decode traces + 2048 chunk trace + 2048 bucket trace at 1,017,275 KV tokens?
Owner

agy continuation (2026-09-25)

1. Consolidation with #38 & Offline Verification

  • Consolidated masked_bucket_trace.py with origin/agy/issue-38; retained branch implementation for strictly safer string parsing and docstrings.
  • Ported bench/fixtures/fixtures-masked-bucket-trace-test.py: 5/5 fixture suites and 18/18 helper unit tests pass offline (23/23 total pass).

2. Simulator De-Risking & Boundaries

  • Ran 2-chip simulated mesh (ttsim-bh-x2, TP=2) with 4-layer Qwen3.5-9B-FP8: initialized FABRIC_1D, 1 GiB trace region, and patched multichip norm is_distributed_norm.
  • Full prefill execution is blocked in simulation by UnsupportedFunctionality: tensix_sfpgt: lreg_dest=9 (ttsim v1.9.9 SFPU limit), matching the documented Stage 2 simulator boundary. Evidence recorded in bench/runs/prefill-bucket-trace-sim-20260925.jsonl. Numerical equivalence requires silicon execution.

3. DRAM Trace-Region Budget

  • Usable P150 DRAM is 30.87 GiB/card. Full 1,017,275 bf8 KV pool (16.66 GiB/card) + weights (13.75 GiB/card) leaves only ~470 MiB headroom [MEASURED, bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl].
  • 4 decode traces + chunk trace consume ~400 MiB TRACE DRAM. Tracing all 5 buckets ([128..2048]) costs ~420 MiB, guaranteeing device OOM [ESTIMATE].
  • Recommendation: Trace Bucket 2048 alone (~140 MiB). Device masks (mask_f32, mask_q, conv_sel) zero out computations past actual_len, dynamically covering any tail 1..2048 tokens without re-tracing.

4. GDN Host Snapshot Analysis

  • In bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl, to_torch/from_torch in _snapshot_gdn_scratch / _restore_gdn_scratch consumes 16.2% of execution samples [MEASURED].
  • Keeping snapshots as host ttnn.Tensors in TILE_LAYOUT avoids untile/retile overhead across 240 tensors (~72 MiB/snapshot), saving ~0.31 s of the 1.94 s tail [ESTIMATE].

5. Hardware Runbook (for execution on cfx-llm2)

  1. Acquire exclusive lock /tmp/ttlock and verify free -m Available > 4 GB.
  2. Baseline warm tail py-spy:
    PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh
  3. E2E benchmark with 2048 bucket trace:
    QWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh
  4. Verify token identity:
    python3 bench/prefix_hit_equivalence.py --model-dir /home/orchestra/tt-models/ThinkingCap-Qwen3.8-27B-FP8
**agy continuation (2026-09-25)** ### 1. Consolidation with #38 & Offline Verification - Consolidated `masked_bucket_trace.py` with `origin/agy/issue-38`; retained branch implementation for strictly safer string parsing and docstrings. - Ported `bench/fixtures/fixtures-masked-bucket-trace-test.py`: 5/5 fixture suites and 18/18 helper unit tests pass offline (23/23 total pass). ### 2. Simulator De-Risking & Boundaries - Ran 2-chip simulated mesh (`ttsim-bh-x2`, TP=2) with 4-layer `Qwen3.5-9B-FP8`: initialized `FABRIC_1D`, 1 GiB trace region, and patched multichip norm `is_distributed_norm`. - Full prefill execution is blocked in simulation by `UnsupportedFunctionality: tensix_sfpgt: lreg_dest=9` (ttsim v1.9.9 SFPU limit), matching the documented Stage 2 simulator boundary. Evidence recorded in `bench/runs/prefill-bucket-trace-sim-20260925.jsonl`. Numerical equivalence requires silicon execution. ### 3. DRAM Trace-Region Budget - Usable P150 DRAM is 30.87 GiB/card. Full 1,017,275 bf8 KV pool (16.66 GiB/card) + weights (13.75 GiB/card) leaves only ~470 MiB headroom [MEASURED, `bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl`]. - 4 decode traces + chunk trace consume ~400 MiB TRACE DRAM. Tracing all 5 buckets ([128..2048]) costs ~420 MiB, guaranteeing device OOM [ESTIMATE]. - Recommendation: Trace **Bucket 2048 alone (~140 MiB)**. Device masks (`mask_f32`, `mask_q`, `conv_sel`) zero out computations past `actual_len`, dynamically covering any tail 1..2048 tokens without re-tracing. ### 4. GDN Host Snapshot Analysis - In `bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl`, `to_torch`/`from_torch` in `_snapshot_gdn_scratch` / `_restore_gdn_scratch` consumes 16.2% of execution samples [MEASURED]. - Keeping snapshots as host `ttnn.Tensor`s in `TILE_LAYOUT` avoids untile/retile overhead across 240 tensors (~72 MiB/snapshot), saving ~0.31 s of the 1.94 s tail [ESTIMATE]. ### 5. Hardware Runbook (for execution on cfx-llm2) 1. Acquire exclusive lock `/tmp/ttlock` and verify `free -m` Available > 4 GB. 2. Baseline warm tail py-spy: `PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh` 3. E2E benchmark with 2048 bucket trace: `QWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.sh` 4. Verify token identity: `python3 bench/prefix_hit_equivalence.py --model-dir /home/orchestra/tt-models/ThinkingCap-Qwen3.8-27B-FP8`
Sign in to join this conversation.
No labels
human-approved
No milestone
No project
No assignees
2 participants
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set

Reference
bitpartner/tt-stack#50
No description provided.