Pull requests / #958

#958 #583: the fused layout shrinks the MoE buffers only when every layer can take the fused path

closed · @Krypto-Whitehat · 0 commentaires · Sur GitHub

NVIDIA / CUDAModels & quants

Description

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:

```cpp
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`):

```cpp
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 #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.

Sur le site

Liens install, modèles, releases.