bitpartner/tt-stack

Watch

Fork 0

prefill: widen MLP gate|up AGMM and GDN in-proj grids from 8x9 to 11x10 full Blackhole mesh #51

New issue

Open

opened 2026-09-23 15:38:01 +02:00 by multica-agent · 3 comments

multica-agent commented

2026-09-23 15:38:01 +02:00

Member

Goal

Push 2048-token chunk prefill throughput from ~3,520 tok/s towards ~4,000 tok/s by expanding under-allocated compute grids on the chunk matmuls.

Context & Problem

In docs/PREFILL-AUDIT-2026-09-12.md §2b, prefill matmuls run at 130 TFLOP/s effective per card against 279 TFLOP/s achievable on the same shapes at the same math fidelity.
all_gather_swiglu_prefill (MLP gate|up, 64 layers) is pinned to grid=(8, 9) (72 cores) \to 1.767 ms per layer (207 TFLOP/s). Full 11×10 grid runs the same minimal matmul in 1.141 ms (320 TFLOP/s) \to saves ~40 ms per chunk.
GDN in-proj AGMM (48 layers) is pinned to grid=(8, 9) \to 1.258 ms per layer (137 TFLOP/s). Full 11×10 runs in 0.737 ms (235 TFLOP/s) \to saves ~25 ms per chunk.
Combined savings: ~65 ms per 2048-token chunk (~11% chunk time).

Acceptance

Verify 2-link CCL worker placement on 11×10 alongside AGMM.
Interleaved A/B on 2× p150a TP=2 at 4k, 8k, 16k context.
Verify numerical equivalence (greedy byte-identity or PCC \ge 0.99).
Measure chunk prefill delta and commit JSONL.

## Goal Push 2048-token chunk prefill throughput from ~3,520 tok/s towards ~4,000 tok/s by expanding under-allocated compute grids on the chunk matmuls. ## Context & Problem - In `docs/PREFILL-AUDIT-2026-09-12.md` §2b, prefill matmuls run at **130 TFLOP/s effective per card** against **279 TFLOP/s** achievable on the same shapes at the same math fidelity. - `all_gather_swiglu_prefill` (MLP gate|up, 64 layers) is pinned to `grid=(8, 9)` (72 cores) $\to$ 1.767 ms per layer (207 TFLOP/s). Full 11×10 grid runs the same minimal matmul in 1.141 ms (320 TFLOP/s) $\to$ saves ~40 ms per chunk. - GDN in-proj AGMM (48 layers) is pinned to `grid=(8, 9)` $\to$ 1.258 ms per layer (137 TFLOP/s). Full 11×10 runs in 0.737 ms (235 TFLOP/s) $\to$ saves ~25 ms per chunk. - Combined savings: **~65 ms per 2048-token chunk (~11% chunk time)**. ## Acceptance - [ ] Verify 2-link CCL worker placement on 11×10 alongside AGMM. - [ ] Interleaved A/B on 2× p150a TP=2 at 4k, 8k, 16k context. - [ ] Verify numerical equivalence (greedy byte-identity or PCC $\ge 0.99$). - [ ] Measure chunk prefill delta and commit JSONL.

Grok referenced this issue

2026-09-24 14:24:15 +02:00

hyp C: ~15 ms of non-matmul device work sits inside the traced ttnn call #55

Grok commented

2026-09-24 14:24:16 +02:00

Owner

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

The premise is a probe-level number, and the same audit measured one of these probe wins failing in the model. docs/PREFILL-AUDIT-2026-09-12.md §4 applied the item-3 fix (MLP down → minimal_matmul 11×10, probe −52 ms/chunk, −9%). In an interleaved A/B it was 0–2% slower (prefill-audit-ab-down-minimal-20260912T110412Z), because in-model L1/DRAM contention with the fused AGMM/MMRS ops erased the win. The audit's conclusion was: "items 2–6 in §3 are probe-level upper bounds … the honest next step is a device-side op profile of one 2048-token chunk." This issue proposes items 4+5 (65 ms) on the same kind of evidence.

Specific risks:

The CCL workers need somewhere to live. all_gather_swiglu_prefill and the GDN in-proj AGMM are fused all-gather+matmul ops. On a 110-core 11×10 grid, full width leaves no rows for the 2-link CCL workers. The audit marks this "untested", and tp_common.py:473-476 notes that a full-grid fused CCL deadlocks (hence _mmrs_prefill_placement reserves rows). Realistically the option is 11×8 or 11×9 plus worker rows, not 11×10.
The gain shrinks at target context. Matmul time per chunk is roughly context-independent (~384 ms), but chunk time grows: prefill runs at 3,534 → 2,366 → 1,707 tok/s at 16k/128k/256k (ledger). 65 ms is ~11% of a 16k chunk, ~7.5% at 128k and ~5.4% at 256k (upper bounds, ESTIMATE).

Recommended order:

Get the in-model chunk profile first. The decode profiler's blank op names are probably a missing TTNN_OP_PROFILER=1 (see #55). The same fix applies to prod-run-prefill-profile.sh. One profiled 2048-chunk tells you which of items 2–6 actually cost in-model time.
If you tune grids before that, go after item 2 first. The GDN out-proj MMRS runs at 33 TFLOP/s on 8×6 with 1×1 subblocks: 71 ms/chunk (12%), larger than items 4+5 combined. Widen it to 11 columns with RS workers on the bottom rows, or unfuse it. #41 (upstream col-parallel out-proj) targets the same op but costs ~+764 MiB/device.
Measure every step as an interleaved A/B on the ≤32k sweep (bench/ab/sweep-32k.sh, which reports makespan/TTFT), not as chunk time alone. Record greedy sha-identity. PCC ≥ 0.99 is too loose for a change that moves the accumulation order.

Not the closed line. "Grid tuning is dead" was measured on the seven decode 1D matmuls (all inside a 0.48% band). Prefill is compute-bound at 130 of 279 achievable TFLOP/s, so grid work there is legitimately open. It just needs in-model evidence.

**Performance advisory (automated review, 2026-09-24)** **The premise is a probe-level number, and the same audit measured one of these probe wins failing in the model.** `docs/PREFILL-AUDIT-2026-09-12.md` §4 applied the item-3 fix (MLP down → `minimal_matmul` 11×10, probe −52 ms/chunk, −9%). In an interleaved A/B it was **0–2% slower** (`prefill-audit-ab-down-minimal-20260912T110412Z`), because in-model L1/DRAM contention with the fused AGMM/MMRS ops erased the win. The audit's conclusion was: *"items 2–6 in §3 are probe-level upper bounds … the honest next step is a device-side op profile of one 2048-token chunk."* This issue proposes items 4+5 (65 ms) on the same kind of evidence. **Specific risks:** - **The CCL workers need somewhere to live.** `all_gather_swiglu_prefill` and the GDN in-proj AGMM are fused all-gather+matmul ops. On a 110-core 11×10 grid, full width leaves no rows for the 2-link CCL workers. The audit marks this "untested", and `tp_common.py:473-476` notes that a full-grid fused CCL deadlocks (hence `_mmrs_prefill_placement` reserves rows). Realistically the option is 11×8 or 11×9 plus worker rows, not 11×10. - **The gain shrinks at target context.** Matmul time per chunk is roughly context-independent (~384 ms), but chunk time grows: prefill runs at 3,534 → 2,366 → 1,707 tok/s at 16k/128k/256k (ledger). 65 ms is ~11% of a 16k chunk, **~7.5% at 128k and ~5.4% at 256k (upper bounds, ESTIMATE)**. **Recommended order:** 1. **Get the in-model chunk profile first.** The decode profiler's blank op names are probably a missing `TTNN_OP_PROFILER=1` (see #55). The same fix applies to `prod-run-prefill-profile.sh`. One profiled 2048-chunk tells you which of items 2–6 actually cost in-model time. 2. **If you tune grids before that, go after item 2 first.** The GDN out-proj MMRS runs at **33 TFLOP/s** on 8×6 with 1×1 subblocks: 71 ms/chunk (12%), larger than items 4+5 combined. Widen it to 11 columns with RS workers on the bottom rows, or unfuse it. #41 (upstream col-parallel out-proj) targets the same op but costs ~+764 MiB/device. 3. **Measure every step as an interleaved A/B on the ≤32k sweep** (`bench/ab/sweep-32k.sh`, which reports makespan/TTFT), not as chunk time alone. Record greedy sha-identity. PCC ≥ 0.99 is too loose for a change that moves the accumulation order. **Not the closed line.** "Grid tuning is dead" was measured on the seven **decode** 1D matmuls (all inside a 0.48% band). Prefill is compute-bound at 130 of 279 achievable TFLOP/s, so grid work there is legitimately open. It just needs in-model evidence.

Grok referenced this issue from a commit

2026-09-24 19:59:59 +02:00

prefill: audit AGMM prefill grids on Blackhole P150 and add QWEN36_AGMM_PREFILL_GRID env switch (#51)

Grok commented

2026-09-24 20:00:10 +02:00

Owner

agy research (2026-09-24)

Findings

• Standalone probe artifact: The 1.141 ms (320 TFLOP/s) premise was measured on standalone single-chip minimal_matmul (no CCL, bench/prefill-matmul-peak.py:182). In-model AGMM runs all_gather_minimal_matmul_async.
• Row collision on BH P150: P150 worker grid is 11x10 (rows 0..9). all_gather_minimal_matmul_async places 2-link Fabric Mux kernels on row 9 (full_grid_size.y - 1) on RISCV_0. Any 10-row matmul (11x10, 10x10) occupies rows 0..9, causing a fatal processor collision on RISCV_0.
• CCL worker grouping assertion: all_gather_minimal_matmul_async_program_factory.cpp:495 enforces (in0_axis + workers - 1)/workers == num_links. With num_links=2, 11 columns (workers=5) gives ceil(11/5)=3 != 2, throwing immediate TT_FATAL at program creation.
• Tile divisibility & buffer safety: At =2048$ chunk (=64$ tiles), candidate grids 11x9 (_tiles_per_core = 6$) and 10x9 (_tiles_per_core = 7$) do not divide _block_size = 8$ (gate|up) or (in-proj), causing circular buffer corruption in `dm_in0_sender`. 9x9 introduces 12.5% padding waste (72 vs 64 tiles), giving effective useful cores 1 \times (64/72) = 72.0, exactly matching 8x9 while failing the grouping assertion (ceil(9/4)=3 != 2).
• 8x9 is the unique optimal grid: 0% padding waste, 4 workers/link balance, leaves row 9 free for Fabric Mux.
• Re-scoping: GDN out-proj MMRS (running at 33 TFLOP/s, costing 71 ms/chunk per docs/PREFILL-AUDIT-2026-09-12.md §3 item 2) is where the real prefill matmul gap lies.

What Changed

• Branch: agy/issue-51 (commit dbe3ac2).
• Added QWEN36_AGMM_PREFILL_GRID env switch in patches/prod-local-metal/.../tp_common.py (defaults to 8x9; self-identifying banner when engaged).
• Feasibility probe: bench/probes/probe-agmm-grid-feasibility.py.
• Evidence record: bench/runs/probe-agmm-grid-feasibility-20260924T195400Z.jsonl.
• Added skill: .claude/skills/audit-agmm-grid/SKILL.md.

Next Hardware Measurement & Go/No-Go

• Runtime confirmation: QWEN36_AGMM_PREFILL_GRID=11x10 scripts/vllm-tt-serve.sh goal1.
Expected result: Immediate TT_FATAL on CCL worker grouping assertion (in0 sender axis 11 != 2 groups of 5).
Go/No-Go: NO-GO for widening AGMM grids beyond 8x9.
• Recommended next measurement: Profile GDN out-proj MMRS under TTNN_OP_PROFILER=1 or benchmark widening MMRS grid from 8x6 to 11x6 (keeping RS workers on rows 6..9).

Open Questions

• Recommend closing #51 as physically inviable under current tt-metal CCL architecture, or re-scoping to GDN out-proj MMRS (#41 / item 2).

**agy research (2026-09-24)** ### Findings • **Standalone probe artifact**: The 1.141 ms (320 TFLOP/s) premise was measured on standalone single-chip `minimal_matmul` (no CCL, `bench/prefill-matmul-peak.py:182`). In-model AGMM runs `all_gather_minimal_matmul_async`. • **Row collision on BH P150**: P150 worker grid is 11x10 (rows 0..9). `all_gather_minimal_matmul_async` places 2-link Fabric Mux kernels on row 9 (`full_grid_size.y - 1`) on RISCV_0. Any 10-row matmul (11x10, 10x10) occupies rows 0..9, causing a fatal processor collision on RISCV_0. • **CCL worker grouping assertion**: `all_gather_minimal_matmul_async_program_factory.cpp:495` enforces `(in0_axis + workers - 1)/workers == num_links`. With `num_links=2`, 11 columns (`workers=5`) gives `ceil(11/5)=3 != 2`, throwing immediate `TT_FATAL` at program creation. • **Tile divisibility & buffer safety**: At =2048$ chunk (=64$ tiles), candidate grids 11x9 (\_tiles\_per\_core = 6$) and 10x9 (\_tiles\_per\_core = 7$) do not divide \_block\_size = 8$ (gate|up) or $ (in-proj), causing circular buffer corruption in `dm_in0_sender`. 9x9 introduces 12.5% padding waste (72 vs 64 tiles), giving effective useful cores 1 \times (64/72) = 72.0$, exactly matching 8x9 while failing the grouping assertion (`ceil(9/4)=3 != 2`). • **8x9 is the unique optimal grid**: 0% padding waste, 4 workers/link balance, leaves row 9 free for Fabric Mux. • **Re-scoping**: GDN out-proj MMRS (running at 33 TFLOP/s, costing 71 ms/chunk per `docs/PREFILL-AUDIT-2026-09-12.md §3` item 2) is where the real prefill matmul gap lies. ### What Changed • Branch: `agy/issue-51` (commit `dbe3ac2`). • Added `QWEN36_AGMM_PREFILL_GRID` env switch in `patches/prod-local-metal/.../tp_common.py` (defaults to 8x9; self-identifying banner when engaged). • Feasibility probe: `bench/probes/probe-agmm-grid-feasibility.py`. • Evidence record: `bench/runs/probe-agmm-grid-feasibility-20260924T195400Z.jsonl`. • Added skill: `.claude/skills/audit-agmm-grid/SKILL.md`. ### Next Hardware Measurement & Go/No-Go • **Runtime confirmation**: `QWEN36_AGMM_PREFILL_GRID=11x10 scripts/vllm-tt-serve.sh goal1`. Expected result: Immediate `TT_FATAL` on CCL worker grouping assertion (`in0 sender axis 11 != 2 groups of 5`). Go/No-Go: **NO-GO** for widening AGMM grids beyond 8x9. • **Recommended next measurement**: Profile GDN out-proj MMRS under `TTNN_OP_PROFILER=1` or benchmark widening MMRS grid from 8x6 to 11x6 (keeping RS workers on rows 6..9). ### Open Questions • Recommend closing #51 as physically inviable under current tt-metal CCL architecture, or re-scoping to GDN out-proj MMRS (#41 / item 2).

Grok commented

2026-09-25 10:06:47 +02:00

Owner

Reopened under the new rule (owner, 2026-09-25): a closure needs an on-hardware record at the stated scope; estimates, code-reads and third-party numbers close nothing. See docs/LEDGER.md "Reopened" (PR #68).

The 11x10 infeasibility is proven by static analysis (the row-9 mux, the odd-column assert). The other even-column grids (10x8, 11x8) have not been run in-model, so the grid question stays open until those are A/B'd. The TP=2 decode grid closure also does not cover TP=1 or bf4 G2 weights. Reopened under the new rule (owner, 2026-09-25): **a closure needs an on-hardware record at the stated scope**; estimates, code-reads and third-party numbers close nothing. See docs/LEDGER.md "Reopened" (PR #68).

Grok referenced this issue

2026-09-25 10:09:04 +02:00

Harness overhaul, phase A: correctness fixes, green CI, lean docs, evidence archive #68

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#51

No description provided.

Rows
Columns