open_kernels: batched prefill for Gemma 3, and the dx_attn hang it uncovered - #101
Merged
Merged
Conversation
The dense block route drove one attention kernel ("dxB") against one
position table ("ptab") for every layer, so a spec with a second layer
type -- Gemma 3's sliding-window dense_local -- couldn't use it.
GemmBlockProgram gains attn_kernel / attn_args (default dxB / ...,
"ptab", so every existing manifest keeps parsing unchanged) and
sandwich / act, so a layer type with its own window names its own
kernel (dxB_local, sharing dxB's instruction stream, patched with the
family's window) and table, and the host chain can run Gemma 3's
sandwich residual norms and GeGLU-tanh alongside the existing plain
SiLU chain.
dense.py's gemm_route() now loops over every layer type instead of
refusing any spec with more than one, and fails closed (no route) for
any family whose (activation, sandwich_norms) pairing the host chain
doesn't cover. Gemma 3-4B and -12B both build to a two-layer-type
route; every other dense family's route is unchanged (schema defaults
verified in manifest_test.cpp).
Unit-tested only: manifest_test.cpp's qwen3 (defaults) and gemma3
(dxB_local, ptab_local, sandwich, gelu_tanh) fixture blocks, a new
mismatched-combo guard test in test_gemma3.py, all 603 spec tests, and
the C++ manifest/block_host/vit tests, including open_qwen36_cli
compiled against the real XRT SDK. No hardware run: building a Gemma 3
kernel set with this route needs the Linux AIE toolchain, which this
box doesn't have, so the argmax/top-5 gate and the decode-after-
block-prefill gate (closing-the-kernel-gap.md #9.6.3 -- the route
writes KV rows sequential decode reads afterward) are still open.
Also archives specs/open-engine/plans/prefill-batch.md: its stage-3
scope (every dense family emitting the route) is now code-complete,
per the plan's own archival trigger. OPEN-PREFILL-BATCH and
OPEN-MANIFEST in spec.md gain the schema clause and a result
paragraph recording what's measured (Qwen3-4B, from the prior
session) and what's still pending (Gemma 3's hardware gate).
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…t NHL < HPO)
dx_attn.py was copied from dx.py before dx.py's og handling changed. It
still put core 0's og on the KV-row fifo -- which forced every og
element to be KVW wide -- and emitted NHL // HPO og elements per core.
For any geometry whose per-core head count is below HPO that count is
ZERO: core 0 never releases an og element, the drain waits forever, and
the first dxB dispatch times out on hardware (ERT state 8). Gemma 3-4B
(NHL 2, HPO 4) and Phi-4-mini (4, 8) are both such geometries; the six
families the route had run on so far all have NHL >= HPO, which is why
it never showed. This ports dx.py's current form: og elements are
kOGH = min(NHL, HPO) heads (attn.h's own constant), N_OG = NHL // kOGH,
and every core drains its own og fifo at Tile(2 + c, 0).
Measured, same box, same sitting, a 1290-token synthetic prompt so the
sliding window is past 1024 rows, block route vs the sequential path:
Gemma3-4B, 34 layers: 8/8 greedy tokens identical after prefill;
prefill 66.8 s vs 86.6 s. 6 layers, 1290 positions: argmax
1274/1290, top-5 1230/1290, min corr 0.99927; 32/32 greedy.
Qwen3-4B (hd 128, the regression check), 6 layers: argmax 1276/1290,
top-5 1205/1290, min corr 0.99989; 32/32 greedy; 10.8 s vs 19.9 s.
Phi4-mini, 4 layers: first block-route run ever; 16/16 greedy;
6.0 s vs 12.1 s.
The flips on both families sit at top-2 margins of a few hundredths of
a logit on a random-id prompt -- the route's bf16 GEMM, same as the
35B's documented behaviour.
Every dense family's dx_attn build changes, so each kernel set needs
re-exporting to pick this up.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Independent confirmation of the
FYI #105 (draft) builds on this branch and touches |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
One or more issues must be addressed before approval.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (2)
What changed in this PR
Adds Gemma 3 batched prefill support and fixes the dx_attn output-fifo hang affecting Gemma 3 and Phi-4-mini.
Changes:
- Adds per-layer attention kernels/tables and Gemma sandwich/GELU handling.
- Corrects attention output partitioning across cores.
- Updates manifests, tests, fixtures, and prefill documentation.
| File | Description |
|---|---|
| src/open_qwen36/manifest.hpp | Updated as part of this pull request. |
| src/open_qwen36/manifest.cpp | Updated as part of this pull request. |
| src/open_qwen36/manifest_test.cpp | Updated as part of this pull request. |
| src/open_qwen36/core.hpp | Updated as part of this pull request. |
| src/open_qwen36/core.cpp | Updated as part of this pull request. |
| specs/open-engine/tests/test_gemma3.py | Updated as part of this pull request. |
| specs/open-engine/tests/fixtures/manifest_qwen3_4b.json | Updated as part of this pull request. |
| specs/open-engine/tests/fixtures/manifest_phi4_mini_4b.json | Updated as part of this pull request. |
| specs/open-engine/tests/fixtures/manifest_hy_mt2_7b.json | Updated as part of this pull request. |
| specs/open-engine/tests/fixtures/manifest_gemma3_4b.json | Updated as part of this pull request. |
| specs/open-engine/spec.md | Updated as part of this pull request. |
| specs/open-engine/plans/archive/prefill-batch.md | Updated as part of this pull request. |
| open_kernels/recipes/dense.py | Updated as part of this pull request. |
| open_kernels/designs/dense/dx_attn.py | Updated as part of this pull request. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
vegah
approved these changes
Sep 23, 2026
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
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.
Closes #99. Closes #100.
Gemma 3 was the one dense family #97 left out. The block route drove one attention kernel (
dxB) against one position table (ptab) for the whole model, and Gemma 3's sliding-window layers need their own — the samedx/dx_localsplit the sequential path already has. Three changes, in three commits:1. The route names its attention kernel per layer type.
gemm_blockgainsattn_kernel/attn_args(defaulting to today'sdxB/ptab, so every existing manifest parses unchanged) plussandwich/act. Gemma 3'sdense_locallayers get adxB_localentry — the samedx_attninstruction stream, patched with the 1024-row window — bound toptab_local, and the host chain learns Gemma's sandwich residual norms and GeGLU-tanh alongside the plain SiLU chain.dense.py'sgemm_route()loops over every layer type instead of refusing anything but one, and fails closed for any family whose activation/norm pairing the host chain doesn't cover.2. A real bug in
dx_attn.py, found the moment Gemma 3 ran. The firstdxBdispatch timed out, reproducibly, sliding window or not.dx_attn.pyhad been copied fromdx.pybeforedx.py's output handling changed: it emittedNHL // HPOoutput elements per attention core, which is zero when a core owns fewer heads than an output element carries — core 0 never releases one, the drain waits forever. Gemma 3 (2 heads per core, 4 per element) and Phi-4-mini (4 and 8) both hit it; the six families the route had run on all haveNHL >= HPO.dx_attn.pynow matchesdx.py(min(NHL, HPO)heads per element, one output fifo per core). Every dense family'sdx_attnbuild key moves, so each kernel set needs re-exporting.3. Spec bookkeeping —
OPEN-PREFILL-BATCHandOPEN-MANIFESTgain the schema clause and two result paragraphs, andplans/prefill-batch.mdis archived: with Gemma 3 in, every dense-shaped family in the catalogue emits the route, which was that plan's last step.Measured
Same box, same sitting, a 1290-token synthetic prompt (long enough that the sliding window is past its 1024 rows), block route vs. the sequential path. The decode check is the one that matters for a sliding-window family: the block route writes the KV rows
dx_localdecodes against afterward, so a mismatch there would corrupt generation while leaving prefill logits looking fine.Every argmax flip on both families is a top-2 margin of a few hundredths of a logit on a random-id prompt — the route's bf16 GEMM, the same behaviour documented for the 35B — and Gemma's profile matches the Qwen3-4B control, so no regression on the hd-128 path. Gemma's 1.30x against Qwen3-4B's 1.66x is expected: past 1024 rows the sequential path's local layers already attend over a capped window, so the route has less to win.
Unit side:
manifest_test.cpp(defaults on the Qwen3 fixture, the two-kernel route on the Gemma fixture), a new mismatched-pairing guard test intest_gemma3.py, all 603 spec tests, andopen_qwen36_clicompiled against the real XRT SDK.Reviewing
Built and run natively on Windows (no WSL) — see the companion docs PR for the toolchain setup, which this repo's docs had wrongly called Linux-only.
🤖 Generated with Claude Code