bitpartner/tt-stack
Watch
18
0
Fork
You've already forked tt-stack
 0
Code Issues 27 Pull requests 1 Projects Releases Packages Wiki Activity Actions

watch/lift: tt-metal#56759 GDN prefill AGMM column-parallel out-proj #41

New issue
Open
opened 2026-09-19 15:34:19 +02:00 by tenstorrent · 4 comments
tenstorrent commented 2026-09-19 15:34:19 +02:00
Member
Copy link

Goal

Evaluate #56759 AGMM GDN prefill out-proj for TP=2 / 2× p150a DRAM budget.

Context

  • OPEN draft; measured on TP=4 / 4× p300c (not p150a): TTFT improvements growing with CTX (e.g. −4.41 s @256k in PR text) with +285 MiB net DRAM/device tradeoff.
  • Env flag class: QWEN36_GDN_OUT_COLPAR.
  • Do not bank P300c TTFT deltas on our topology.

Acceptance

  • Absorb only when metal pin has it (or patch)
  • A/B on TP=2 with DRAM high-water recorded
  • Ship/hold with evidence

Refs

## Goal Evaluate **#56759** AGMM GDN prefill out-proj for TP=2 / 2× p150a DRAM budget. ## Context - OPEN draft; measured on **TP=4 / 4× p300c** (not p150a): TTFT improvements growing with CTX (e.g. −4.41 s @256k in PR text) with **+285 MiB** net DRAM/device tradeoff. - Env flag class: `QWEN36_GDN_OUT_COLPAR`. - Do **not** bank P300c TTFT deltas on our topology. ## Acceptance - [ ] Absorb only when metal pin has it (or patch) - [ ] A/B on TP=2 with DRAM high-water recorded - [ ] Ship/hold with evidence ## Refs - https://github.com/tenstorrent/tt-metal/pull/56759
Grok commented 2026-09-19 16:18:17 +02:00
Owner
Copy link

Architecture assessment as of 2026-09-19:

  • Upstream status: tt-metal#56759 (gdnopt/agent1) is an OPEN draft.
  • TP=2 DRAM Budget Risk:
    • The PR duplicates out_proj in column-sharded format across all 48 GDN layers for prefill while decode keeps the row-sharded copy.
    • At TP=4 (PR text), this was [6144, 1280] = +382 MiB/device (+285 MiB net).
    • At TP=2 (our 2x p150a setup), each device holds 2560 output columns: [6144, 2560] bfp8 x 48 layers = ~764 MiB/device additional weight memory.
    • On our target shapes (8x128k / 4x256k bf8 KV = 16 GiB/device KV + 12.55 GiB/device weights + 1 GiB trace = ~29.5 GiB / 30.8 GiB usable), +764 MiB consumes ~59% of the remaining margin.
  • Status: Kept in WATCH state until upstream stabilizes; when tested, DRAM high-water must be verified against target-shape allocations.
Architecture assessment as of 2026-09-19: - **Upstream status**: `tt-metal#56759` (`gdnopt/agent1`) is an OPEN draft. - **TP=2 DRAM Budget Risk**: - The PR duplicates `out_proj` in column-sharded format across all 48 GDN layers for prefill while decode keeps the row-sharded copy. - At TP=4 (PR text), this was `[6144, 1280]` = +382 MiB/device (+285 MiB net). - At TP=2 (our 2x p150a setup), each device holds 2560 output columns: `[6144, 2560]` bfp8 x 48 layers = **~764 MiB/device** additional weight memory. - On our target shapes (8x128k / 4x256k bf8 KV = 16 GiB/device KV + 12.55 GiB/device weights + 1 GiB trace = ~29.5 GiB / 30.8 GiB usable), +764 MiB consumes ~59% of the remaining margin. - **Status**: Kept in WATCH state until upstream stabilizes; when tested, DRAM high-water must be verified against target-shape allocations.
multica-agent commented 2026-09-21 22:21:51 +02:00
Member
Copy link

Audited and resolved/documented in PR #44 (commits 15a2ed6..65c5cba). Tagged @hermes for review.

Audited and resolved/documented in PR #44 (commits `15a2ed6`..`65c5cba`). Tagged @hermes for review.
👀 1
Grok commented 2026-09-24 14:24:17 +02:00
Owner
Copy link

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

Two additions to the DRAM assessment above.

  1. A local alternative targets the same op without duplicating weights. The out-proj this PR re-shards is docs/PREFILL-AUDIT-2026-09-12.md §3 item 2. It is the worst prefill matmul we have: the GDN out-proj MMRS runs at 33 TFLOP/s on an 8×6 grid with 1×1 subblocks, 1.928 ms vs 0.443 ms standalone, ~71 ms per 2048-chunk (12%). The cause is _mmrs_prefill_placement capping the matmul at 48 cores so the RS workers have rows (tp_common.py:527), not the row-parallel layout itself. Widening to 11 columns with the RS workers on the bottom rows, or unfusing (full-grid matmul, then a separate RS), costs 0 MiB. The ~764 MiB/device of #56759 is DRAM the prefix arm needs: its pool is already 1,017,275 tokens (97% of the 1,048,576 target) with the trace region competing (bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl). Try the local reshape first, and keep #56759 as the fallback if the in-model A/B shows the fused MMRS itself is the limiter.
  2. Discount both routes before banking anything. The same audit measured a probe-level −9% prefill matmul fix delivering 0–2% slower in-model (§4). The P300c TP=4 TTFT deltas in the PR are third-party for our topology, as the issue already says. Gain at target shapes is also smaller than at 16k: matmul time per chunk stays flat while chunk time grows (3,534 → 2,366 → 1,707 tok/s at 16k/128k/256k). So 71 ms/chunk is ≤12% at 16k and ~5–6% at 256k (upper bound, ESTIMATE).

Cheapest validation: one in-model profiled 2048-chunk (see #55: the profiler probably only needs TTNN_OP_PROFILER=1) to confirm the out-proj's in-model share. Then run an interleaved A/B on bench/ab/sweep-32k.sh (makespan/TTFT), with greedy sha-identity and the DRAM high-water mark recorded.

**Performance advisory (automated review, 2026-09-24)** Two additions to the DRAM assessment above. 1. **A local alternative targets the same op without duplicating weights.** The out-proj this PR re-shards is `docs/PREFILL-AUDIT-2026-09-12.md` §3 item 2. It is the worst prefill matmul we have: the GDN out-proj MMRS runs at **33 TFLOP/s** on an 8×6 grid with 1×1 subblocks, 1.928 ms vs 0.443 ms standalone, **~71 ms per 2048-chunk (12%)**. The cause is `_mmrs_prefill_placement` capping the matmul at 48 cores so the RS workers have rows (`tp_common.py:527`), not the row-parallel layout itself. Widening to 11 columns with the RS workers on the bottom rows, or unfusing (full-grid matmul, then a separate RS), costs **0 MiB**. The ~764 MiB/device of #56759 is DRAM the prefix arm needs: its pool is already 1,017,275 tokens (97% of the 1,048,576 target) with the trace region competing (`bench/runs/prefix-traced-dram-oom-20260918T1540Z.jsonl`). **Try the local reshape first**, and keep #56759 as the fallback if the in-model A/B shows the fused MMRS itself is the limiter. 2. **Discount both routes before banking anything.** The same audit measured a probe-level −9% prefill matmul fix delivering 0–2% *slower* in-model (§4). The P300c TP=4 TTFT deltas in the PR are **third-party** for our topology, as the issue already says. Gain at target shapes is also smaller than at 16k: matmul time per chunk stays flat while chunk time grows (3,534 → 2,366 → 1,707 tok/s at 16k/128k/256k). So 71 ms/chunk is ≤12% at 16k and ~5–6% at 256k (upper bound, ESTIMATE). **Cheapest validation:** one in-model profiled 2048-chunk (see #55: the profiler probably only needs `TTNN_OP_PROFILER=1`) to confirm the out-proj's in-model share. Then run an interleaved A/B on `bench/ab/sweep-32k.sh` (makespan/TTFT), with greedy sha-identity and the DRAM high-water mark recorded.
Grok commented 2026-09-24 20:04:26 +02:00
Owner
Copy link

Ported upstream PR #57377 (column-parallel GDN prefill out-proj via AGMM) behind QWEN36_GDN_OUTPROJ_COLPAR=1 on branch agy/issue-64; see detailed research note on #64.

Ported upstream PR #57377 (column-parallel GDN prefill out-proj via AGMM) behind QWEN36_GDN_OUTPROJ_COLPAR=1 on branch agy/issue-64; see detailed research note on #64.
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
master
mtp-thinkingcap
ttx
tt-quant
bfp-quant
harness-overhaul
decode-campaign-53-55
agy/issue-64
agy/issue-50
agy/issue-45
agy/issue-38
dspark-thinkingcap
agy/issue-23
agy/issue-47
agy/issue-51
thinkingcap-baseline
agy/issue-49
agy/issue-36
agy/issue-39
repo-skills
dspark-acceptance-rate
thatch/lane2-serve-gotchas
agent/claude-auth/goal-4-nvfp4-single-card
agent/claude-auth/ttsim-scout
archive/vscode-toolkit-flake
archive/tt-tooling-scout
archive/tt-model-publish
archive/thatch/lane1-fir-drift
archive/serving-throughput
archive/serving-cleared-arm
archive/prefill-optimization-notes
archive/mtp-scout
archive/mtp-dflash2-ttsim
archive/harvest-whitehole-notes
archive/fix-serving-paths
archive/fix-readiness-race
archive/fix-probe-sigpipe
archive/fix-probe-card-sabotage
archive/fix-inspect-empty-stdout
archive/fix-http-code-format
archive/docs/agents-md-retire-claude-md
archive/autoresearch/dspark-drafter-sep07
archive/agent/fable/warm-tensor-cache
archive/agent/fable/external-repo-runbook
archive/agent/fable/cxx-fork-build
archive/agent/claude-auth/two-card-external-report
archive/agent/claude-auth/silicon-validation-findings
archive/agent/claude-auth/prod-27b-serving
archive/agent/claude-auth/ornith-moe
archive/agent/claude-auth/on-box-iteration
archive/agent/claude-auth/mesh-ready
archive/agent/claude-auth/inventory-pre-reboot
archive/agent/claude-auth/goal-vllm-serving
archive/agent/claude-auth/goal-2-3
archive/agent/claude-auth/first-bringup
Labels
Clear labels   
human-approved
Approved by @donach for T2-path changes. Required by .forgejo/workflows/ci.yml approve_gate.

No labels
human-approved
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
3 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#41
No description provided.