改动 / #2903
#2903 metal: the gate/up grid follows the tile, so a K-quant layer covers its whole ffn width
open · @ctr-137 · 0 评论 · 去 GitHub 看
Server & APIModels & quantsDocumentationWindows
说明
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)
本站相关内容
相关页面的快捷入口。