prefill: capture and replay masked-bucket traces (<=2048 tok) to eliminate ~420ms host dispatch floor #50
Labels
No labels
human-approved
No milestone
No project
No assignees
2 participants
Notifications
Due date
No due date set.
Dependencies
No dependencies set
Reference
bitpartner/tt-stack#50
Loading…
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Goal
Enable trace replay for prefill buckets
\le 2048tokens 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
patches/prod-local-metal/models/demos/blackhole/qwen36/tt/model.py:2580-2615, any prompt whereactual_len < chunk_size(2048) routes toprefill_masked_bucket.prefill_masked_bucketruns all 64 layers through_forward_prefill_chunk_maskedeagerly without trace execution.docs/PREFILL-AUDIT-2026-09-12.md: eager dispatch adds a flat ~420 ms host CPU dispatch floor from sequential Python / TTNN op dispatch:\to491 ms (1,106 tok/s).\to1,006 ms (2,044 tok/s).Mechanism
capture_prefill_trace_bucketmachinery (model.py:1649) for fixed bucket sizes [128, 512, 1024, 2048].ttnn.execute_traceduringprefill_masked_bucket, collapsing host dispatch to <5 ms.Acceptance
\ge 0.99).\le 2048tokens exceeds 3,000 tok/s.bench/runs/.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:
prefill-cold-pyspyharness. 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_scratchrunto_torch/from_torchover 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 hostttnntensors is a smaller change than this issue and should go first.actual_len. Do not start with four buckets: each captured trace consumes trace-region DRAM, andbench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonlshows the trace region already competing with the KV pool. Steal before you write: per the #38 inventory, tt-metal#56119 carriesQWEN36_PREFILL_BUCKET_TRACE=1.Correctness traps specific to a fixed-bucket trace:
prefill_traced_chunkeddocstring 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 forcebeta/gto zero pastactual_lenfrom a runtime device tensor, thecommitted_masktrick noted in #22. Otherwise the GDN state after the tail is wrong, and the hit's first decode token is wrong too.model.py:2376, 2451, 2588). Capture during warmup only.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.
agy research (2026-09-24)
Findings & Architectural Design
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].prefill_masked_bucketadds ~420 ms flat dispatch floor across 64 layers (~151 hostfrom_torchuploads) [MEASURED, docs/PREFILL-AUDIT-2026-09-12.md]. Cold prefix py-spy showed 16.2% host tensor marshalling (to_torch/from_torchin_snapshot_gdn_scratch/_restore_gdn_scratch) [MEASURED, bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl].2048rather 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].mask_f32zeroes\betaand decaygfort \ge actual\_len, halting DeltaNet state updates;conv_selone-hot matmul extracts the realK-1inputs. Fixed-widthfill_ptdiverts pad tokens to a scratch KV block (QWEN36_PREFILL_BUCKET_EXTRA_BLOCK=1), keeping block 0 safe._forward_prefill_bucket_body_tpeager 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 intests/test_masked_bucket_trace_helpers.py.layer.py&gdn/tp.py: Threadgdn_masks, replacing dynamic FIR slicing with one-hot matmulconv_sel.model.py&qwen36_vllm.py: Capture and replay masked bucket traces behind env switchQWEN36_PREFILL_BUCKET_TRACE(default OFF); allocate +1 pad KV block.modules/tenstorrent-serving.nix: Bind mounts, assertion, and self-identifying banner checks forprefillBucketTrace.bench/ab/prefill-warm-pyspy-prefix.sh& skill: Reusable py-spy profiler for warm prefix-cache tails.Next Hardware Measurements (on cfx-llm2)
PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.shQWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.shbench/prefix_hit_equivalence.py.Open Questions
to_torch/from_torchentirely?agy continuation (2026-09-25)
1. Consolidation with #38 & Offline Verification
masked_bucket_trace.pywithorigin/agy/issue-38; retained branch implementation for strictly safer string parsing and docstrings.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
ttsim-bh-x2, TP=2) with 4-layerQwen3.5-9B-FP8: initializedFABRIC_1D, 1 GiB trace region, and patched multichip normis_distributed_norm.UnsupportedFunctionality: tensix_sfpgt: lreg_dest=9(ttsim v1.9.9 SFPU limit), matching the documented Stage 2 simulator boundary. Evidence recorded inbench/runs/prefill-bucket-trace-sim-20260925.jsonl. Numerical equivalence requires silicon execution.3. DRAM Trace-Region Budget
bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl].mask_f32,mask_q,conv_sel) zero out computations pastactual_len, dynamically covering any tail 1..2048 tokens without re-tracing.4. GDN Host Snapshot Analysis
bench/runs/prefill-cold-pyspy-prefix-traced-20260917T193539Z.jsonl,to_torch/from_torchin_snapshot_gdn_scratch/_restore_gdn_scratchconsumes 16.2% of execution samples [MEASURED].ttnn.Tensors inTILE_LAYOUTavoids 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)
/tmp/ttlockand verifyfree -mAvailable > 4 GB.PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.shQWEN36_PREFILL_BUCKET_TRACE=2048 PATCHED_DIR=/tmp/warmpyspy/sources bench/ab/prefill-warm-pyspy-prefix.shpython3 bench/prefix_hit_equivalence.py --model-dir /home/orchestra/tt-models/ThinkingCap-Qwen3.8-27B-FP8