Pull requests / #2903

#2903 metal: the gate/up grid follows the tile, so a K-quant layer covers its whole ffn width

open · @ctr-137 · 0 Kommentare · Auf GitHub

Server & APIModels & quantsDocumentationWindows

Beschreibung

The prompt chain's gate/up GEMM is dispatched over the wrong grid on a Mac without the M5 right-input layout, for
any pack whose routed gate/up is a K-quant.  It computes 160 of each expert's 640 gate/up rows and leaves the rest as
the previous layer left them, and the down GEMM reads that: wrong tokens, no error, no warning.

### What happens

The prompt GEMMs are dispatched indirectly, and `mm_scan` publishes the grids one MoE layer at a time.  Every one of
those fields is per layer — `meta[3]` is this layer's tile count — except the gate/up grid x, which the host wrote
**once per chunk** as the constant `NFF / 64` (`window.cpp:2132`, `prompt_chain.inc:96`).

The tile is not a constant.  It is chosen from the layer's quant type, and `gu2_tile()` hands Q4_K, Q5_K and Q6_K a
**TN-16** tile (`_16_16_128_2`).  `mm_gu2` writes TN output columns of `HS` per threadgroup (`moe_mm.metal:302`,
`f0 = tg.x * TN`), so the launch has to cover `NFF / TN` threadgroups.  With `NFF` 640 the constant gave 10 and the
K-quant tile needs 40: rows 0..159 written, 160..639 stale.

The kernel's own header said so — `moe_mm.metal:47`, *"gx_gu in meta[2] ... written by the host once"* — and the
dispatch picks the tile the grid has to match:

```cpp
} else {
    const std::string fn = "mm_gu2_" + std::to_string(eg.type) + gu2_tile(eg.type);   // TN 16 for the K-quants
    gu.tg[0] = 64;
}
```

### Reachability

No flag needed.  `docs/MAC.md:61` offers `--model ...|UD-Q6_K_XL` on any Apple Silicon Mac, and
`docs/UNSLOTH_Q6.md:35-38` records that pack's `ffn_gate_exps`/`ffn_up_exps` as Q6_K in 47 of 48 layers, with
`ffn_down_exps` Q8_0 (whose `mm_down2` instance exists, so the whole prompt path runs).  On an M1–M4 every prompt
chunk of 64 rows or more then multiplies garbage into the hidden states.  The non-M5 configurations the docs
measured use TN-64 tiles, which is why it had not been seen.

### What I changed

- **`metal/src/moe_tiles.h` (new)** — the v2 tiles and the TN each one writes, returned together so they cannot
  drift apart, plus the v3 tiles' TN/SG and the grid rule.
- **`metal/kernels/moe_mm.metal`** — `MmPlanArgs` carries `gu_gx`/`dn_gx` and `mm_scan` writes `meta[2]`/`meta[5]`
  from it.  It already writes the sibling fields per layer; one host value for the whole chunk cannot serve a tile
  that varies by layer.
- **`window.cpp` / `prompt_chain.inc`** — the plan carries the grid of *this layer's* tile, the v3 arm derives it
  from the tile it names, and the host writes are gone.
- The decode window needed nothing: it already derives its grid from the tile it picks
  (`window.cpp:834`, `NFF / (gs * (32 / gl) * gr)`).

### Why no test caught it

Both harnesses picked their v2 tile from a rule of their own instead of the engine's, so the engine's TN-16 dispatch
was never in front of a parity run:

```cpp
// prompt_chain.cpp, v2: TN 64 for every type
(fmt_unit32(eg.type) ? "_16_64_32_2" : "_16_64_16_2")
```

Both now use `gu2_tile`/`dn2_tile`, and `moe_mm_parity`'s v2 default parses the tile's own name back and refuses to
run if the name and its TN disagree.

### Tests

- **`metal/tests/moe_tile_grid_test.cpp` (new)** — a host test: no Metal device, no model, links nothing, registered
  with ctest, so it runs on a configured tree whether or not there is a card.  It checks the tile table, that every
  tile's name parses back to the TN it declares, that the TN divides the output width (an extra threadgroup writes
  past `HS`, a missing one leaves the tail unwritten), and that a K-quant layer's grid is 40 where the constant gave
  10.
  Compiled and run with MSVC 19.51 `/std:c++20`: **39 checks, all pass.**  With the header's K-quant TN put back to
  64 as a counter-proof, **four fail and it exits 1**.
- **`moe_mm_parity.cpp` / `prompt_chain.cpp`** — updated for the plan struct and to run the engine's own tiles.

### What I could not check

**Nothing here was built or run on macOS.** This review was done on Windows with no GPU and no Metal toolchain, so
the two harness changes are *not* run: they are the runs that would confirm the fix end to end, and they are the kind
of thing `docs/TEST_REQUESTS.md` asks for.  What I have is the host test above and the reading in this description.

### What I left alone

- The v3 (M5) arms already derived their grids from the tile they dispatch, in both files.
- `dn2_tile` returns TN 64 for every type, so the down grid was never wrong.  It is derived the same way now for
  symmetry, not because it was broken.
- The default path stays byte-identical: the v3 arm and every TN-64 tile compute the same grids as before.  The only
  layers whose dispatch changes are the v2 K-quant ones, which were wrong.

There is no issue for this — it came out of a static audit of the tree rather than a report.  Happy to file one if
you would rather triage it there first.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Mehr auf der Site

Links zu Install, Modellen, Releases.