[Krypto-Whitehat]

Fix for #954.

The defect

fused_layout() shrinks a layout's GU, H and Xq to the fused path's sizes, with an MMQ-sized floor for stream_all_min() - 1 = 1023 tokens, so that a last sub-1024-token chunk can still run MMQ in them. It decided to do that on:

return src && fused_ring() && mmq_plan().any && ring_slots(T) > STAGE && (int64_t) T >= stream_all_min();
//                                                 ^^^^^^^^^^^^^ ONE layer's experts go through MMQ

The run loop picks the path per layer (prefill.cpp:2286):

const bool use_mmq = mmq_plan().any && mmq_plan().layer[(size_t) l];

A layer MMQ does not cover keeps the FP16 path, and that path writes the layer's routed rows of the full chunk into exactly those buffers — m.gemm.f16(..., m.GU + o0 * 1280, ne, 1280, N) over ne rows (prefill.cpp:2804), and MMQ likewise at T * K rows into Xq (prefill.cpp:2428). The fused layout holds 1023 tokens of those rows, so at a 16384-token chunk the write runs ~16x past the end: garbage, an illegal access, or a kernel that never returns. The first ck() to reach the sticky error is prefill mmq: iota a prompt later, which is why the report reads as a hang in the ring rather than as a buffer that was too small.

This is the same hole fused_ring() closes one level up for the native kernels (220e0e8): that one asks "do the fused kernels cover every layer", this asks "do MMQ's", and neither covered both. The shrink is now all-or-nothing over both, and the SYCL mirror matches.

Why the ring shows up in the report

The byte-budget ring is not at fault. It moved the chunk (16384 on this rig), and the chunk is what sets how far past the buffer the write runs. That is also consistent with all three control runs on the same binary:

env	chunk	result
STRATA_PF_FUSED=1 (default ring)	16384	overflow → hang / prefill mmq: iota
STRATA_PF_FUSED=0 (default ring)	16384	fused off, so fused_layout() is false → MMQ-sized buffers → works
STRATA_PF_FUSED=1 + STRATA_RING_BYTES=0	0.1.39's list	same shrink, smaller chunk → works

It is also consistent with 0.1.37 running the same configuration for ~4,700 requests: its chunk was below the fused layout's 1023-token MMQ floor, or the layer that overflows was not on the prompt path.

What this PR changes

1. mmq_all() — a new predicate next to mmq_plan(); fused_layout() requires it instead of mmq_plan().any. A pack with one layer MMQ does not cover now keeps MMQ's buffers, as fused_ring() already does for the native kernels.
2. m.fused_bufs records which layout carve() made, and the run loop refuses a layer that would write full-chunk rows into the fused buffers — with the layer and the chunk named, instead of hanging. This is the backstop for anything that drifts between the two decisions later.

On a pack where every layer is both MMQ-capable and natively covered (including a homogeneous IQ3_XXS one), mmq_all() == mmq_plan().any and the fix changes nothing: no chunk, no loan, no buffer size. On a pack with an odd layer out, the fused layout switches off there, exactly as it already does for the native kernel coverage.

Verification — please read, this is not runtime-verified

I do not have a CUDA toolchain or a Blackwell card on the machine this was prepared on, so I could not build strata_prefill, run the test suite, or reproduce #954. What I did verify:

• The new predicate compiles clean under MSVC /W4 and returns the right answer for every layer vector, exercised in a standalone harness with mmq_all() copied verbatim from the patch: all-MMQ → shrink; one FP16 layer (first, middle, last) → no shrink; no MMQ / no layers → no shrink. That harness confirms the hole directly — with 48 layers and serve: fix zero-token replies from engine-queue race between requests #7 on the FP16 path, the old predicate says shrink and the new one says no.
• The existing prefill_fused_moe_test / prefill_fused_iq_test do not cover this: they exercise the kernels with their own buffers and never build a layout, so the ring × fused × layout interaction has no test. I did not add one either — mmq_all() is inline in an anonymous namespace over a loaded pack, so a test would have to duplicate it rather than call it, which would test the copy.

So the overflow mechanism above is established by reading the code, and the guard will name the offending layer on the first run if some other condition is what actually triggers #954 on that rig. The reporter should confirm against their own 0.1.39 build — a 0.1.39 binary with only this patch, same rig. If it still hangs, the new message will say which layer and which chunk, and that is the next thread to pull. STRATA_RING_BYTES=0 remains a valid workaround meanwhile.

Not changed on purpose

bytes_needed() prices the loan with fused_layout(T, true) while carve() uses the real m.src != nullptr. When a layout is created without an expert source the loan is under-priced, and Alloc::take bounds-checks, so the result is a clean "device buffers ... do not fit" at init rather than a hang. Loud, not silent, and out of scope here — happy to fix it in a follow-up if you want the two call sites to share one predicate.

[Niko1221]

Thanks for the careful read. The overflow you describe was real on 0.1.39: fused_ring() was true for a native pack, but a layer that MMQ had no tile for (mmq_plan().layer[l] == 0) took neither MMQ nor the fused kernels at a streamed chunk, so it wrote the full chunk's rows into the shrunk GU/H/Xq. In 0.1.40 the run loop already sends such a layer to the fused kernels (73a1a4d: layers MMQ does not cover but the native kernels do are marked fused-only, and fused_ring() already requires the native kernels to cover every layer). Every layer of a pack that gets the shrunk layout therefore runs fused or MMQ at a streamed chunk; none reaches the FP16 path. We are not merging mmq_all(): in 0.1.40 it would switch the shrink off for those fused-only layers (for example UD-Q4_K_XL on gfx11) and give back the VRAM the fused layout saves, with no overflow left to prevent. Closing this one; the guard idea is a good one and we may add it separately.

[Niko1221]

Thank you, Krypto-Whitehat. The root cause you found is fixed in 0.1.40 (fused_ring and the per-layer plan), and the backstop error from your PR is in 0.1.40.2 with you as co-author.