Pull requests / #2652

#2652 Metal: the routing trace of a resident window (STRATA_METAL_DUMP_ROUTING)

open · @alvarorsouza-arch · 0 Kommentare · Auf GitHub

BenchmarksAMD / HIPNVIDIA / CUDAModels & quants

Beschreibung

## What this changes

`--dump-routing` and `tools/collect_profile_traces.py --mode routing` work on the pooled backends because the routed ids reach the host through the doorbell the expert pool reads (`src/program/generate.cpp:5929` prints exactly that when you pass `--no-pool`). The Metal engine has no host pool, it reads experts in place from the mapped file, so on a Mac the flag is accepted and the trace comes out empty:

```text
routing dumped           /tmp/one.trace (0 records of layer, k, ids, weights)
verify window            wait for rings 0.000  pool 0.000  host 0.000  commit 0.000 ms/round;
                         CPU experts 0.00 distinct / 0.00 routed per layer
```

The file exists, so `tools/collect_profile_traces.py` reports success, and `tools/make_profile.py` has nothing to read. That blocks the input `tools/quant_bound.py` wants from the engine's window logs, and falls back to `--union` (default 1.0) without them: U(P), the distinct experts a pass of P tokens reads per expert a single token reads. The window logs give it on the pooled backends because they count distinct against routed experts per layer; on Metal that line reads `CPU experts 0.00 distinct / 0.00 routed per layer`.

Two files, 66 added lines, Metal only:

- `metal/kernels/moe.metal`: a `route_dump` kernel that copies one layer's `moe_route` output into that layer's own slot. It has to be a device copy because `ids_` is reused by the next layer and a resident window never hands the host a turn between layers.
- `metal/src/engine/window.cpp`: `STRATA_METAL_DUMP_ROUTING=path` opens the file, allocates one slot pair per layer, dispatches `route_dump` right after the route dispatch in `Window::layer` (the main model's window builder only, not the draft one, since both would write slot 0), and `finish()` writes the window's records in the format `tools/make_profile.py:read_trace` reads: `int32 layer, int32 k, k int32 ids, k float weights`.

## What it leaves alone

Opt-in by environment variable, default path untouched. With the variable unset no buffer is allocated, no dispatch is added to the window, `routing_write()` returns on its first line, and no file is opened. The core is unchanged. It is an environment variable rather than a flag because the core exposes no routing trace path to the backend (nothing in `include/` or `src/core/` carries one); it follows `STRATA_METAL_KPROF` next to it. Tell me where you would rather hang it and I will move it.

CUDA, HIP and SYCL builds are not affected: `metal/` is its own CMake project and does not enter those builds, and no shared file is touched.

## How I checked

Apple M3 Max, 40-core GPU, 128 GB unified memory, macOS 26.6, in high power mode. Three one-shot runs of one command (`--spec 2` with the MTP draft layer, `--expert-cache 64`, `--max-context 65536`, experts resident, `--tokens-file` a 37 token prompt, `--max-new 120 --greedy`):

- a build of `main` at `296b59d` with no patch: 61.94 tok/s
- the patched build, variable unset: 61.21 tok/s
- the patched build, variable set: 60.97 tok/s

The `output  :` line is the same 120 token ids in all three, byte for byte.

The trace is 519,552 bytes: 48 layers by 123 window rows by 88 bytes. `tools/make_profile.py:read_trace` reads it as is, and walking the file gives 5,904 records, all with layer in 0 to 47, k equal to 10 and ids in 0 to 511, no malformed record and no trailing bytes, with weights from 0.0084 to 0.8066.

The cost is small but I cannot pin it down: this pair says 0.4% (60.97 against 61.21), an earlier pair on the same build said 2.9% (61.31 against 63.11), and pure decode on this machine moves about 7% from boot to boot. One extra dispatch per layer per window is the whole overhead, so it should stay under a few percent.

## What it bought

U(P) over 48 layers of a run like that, greedy, one prompt, counted as the distinct experts in the pass divided by the distinct experts of one token in it:

| Pass | U(P) mean | across layers | experts repeated in the pass |
| --- |---:|---|---:|
| 2 tokens | 1.684 | 1.543 to 1.950 | 15.8% |
| 4 tokens | 2.822 | 2.288 to 3.765 | 29.5% |
| 8 tokens | 4.615 | 3.301 to 6.773 | 42.3% |

Worth one caveat about the rows: they are the window's proposals, so a draft that was rejected still shows the experts it routed. For a cache profile that is the useful view; for accepted token accounting you would filter them.

The machine and the unmodified binary numbers are in the report PR #2649. Related but different: #1547 puts routing counts on `/metrics` in the server, which is not what reads a per window trace, and nothing open covers the empty trace on Metal. If you would rather have `collect_profile_traces.py` fail on a 0 record trace than accept it, that is a one line change I can add here or in an issue.

Mehr auf der Site

Links zu Install, Modellen, Releases.