Skip to content

moe_placement: Kernighan-Lin refinement, and the held-out share next to the in-sample one (Refs #1702) - #1947

Merged
bong-water-water-bong merged 1 commit into
mainfrom
cl-infra-1702-kl
Nov 10, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
cl-infra-1702-kl

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1702. Two findings from running the merged tools/moe_placement.py on the report's own traces (~/wt/megacoa/*.tsv on strixhalo, read-only), and the fixes for both.

1. The merged tool does not reach the report's numbers

origin/main (d9d80c27), tools/moe_placement.py <trace> -g 2 -e <experts>:

trace report (#1702 §1) merged tool this PR
Qwen3.6 chat (28 800 decode steps) 46.2 % 34.8 % 51.1 %
Qwen3.6 code (8 000) 60.0 % 34.8 % 55.6 %
Flash-Next (10 560) 44.0 % 24.7 % 52.3 %
contiguous control 0.6 / 0.5 / 0.2 % 0.6 / 0.5 / 0.2 % same

The controls match the report exactly, so the metric is the same one. The difference is the partitioner: the tool cut a greedy max-affinity chain into G runs and its docstring said that "captures the clustering" without the Kernighan-Lin pass the report used. This PR adds the KL refinement (pairwise between groups, equal-size swaps, so a balanced split stays balanced). Each trace runs in under 2 s. G=4 on chat gives 27.0 % (the report: 25.3 %).

2. All of those shares are in-sample, and held out they are much lower

Every number above, the report's included, scores a table on the same decode steps it was built from. A placement is fixed in the model while prompts change, so I measured what a table does on steps it never saw:

table built from scored on Qwen3.6 chat Qwen3.6 code Flash-Next
whole trace same trace 51.1 % 55.6 % 52.3 %
first half of each layer's steps second half 22.5 % 26.3 % 18.9 %
the other prompt's trace this trace 7.0 % 13.5 % not measured

A table fitted to chat and code together scores 50 % on both in-sample, which is what overfitting looks like at a few hundred steps against 256 experts. The placement is still 15 to 50 times the contiguous control on unseen steps, but it is a 7–26 % lever, not 44–60 %. That matters for the gate both decision comments on #1702 set for the consumer.

The tool now prints both shares, writes them into the table (share_in_sample, share_held_out, held_out_steps; additive keys, schema name unchanged), and has --score TABLE to score an existing table on another trace. docs/moe-streaming.md carries the table above and says which share is which.

Proof

$ python3 tools/moe_placement.py ~/wt/megacoa/qwen36-chat.tsv -g 2 -e 256 -o chat.json
moe_placement: 40 layer(s), 2 group(s); one-device share 51.1% over 28800 decode steps IN-SAMPLE (contiguous control 0.6%); held out 22.5% over 14400 steps (table from the first half of each layer's steps, scored on the second half)
$ python3 tools/moe_placement.py ~/wt/megacoa/qwen36-code.tsv --score chat.json
moe_placement: chat.json on .../qwen36-code.tsv: one-device share 13.5% over 8000 decode steps (contiguous control 0.5%)
$ python3 tools/moe_placement.py ~/wt/megacoa/fn-dec.tsv --score chat.json
moe_placement: chat.json has no layer 40 (8 of the trace's 48 layer(s) are not in the table): it was built for a different model or a shorter trace     rc=2

Self-test (ctest moe_placement): 21 ok, 0 failures. The new bridge-pair arm against origin/main's build():

groups [[0, 1, 3, 5], [7, 2, 4, 6]] share 37.5% of 160  -> the new arm wants 62.5%: FAIL

Other new arms: group sizes unchanged by the refinement; on a structureless random trace the refined plan holds at least the chain's co-activation weight; a table from the first half of a layer whose clusters then change scores 0 % held out, and 100 % when they do not; --score refuses a file with another schema.

Not tested, not changed

  • No consumer and no end-to-end rate. Nothing in the engine reads the table; this is still an offline artifact. The NVMe write order stays gated on the HRX streamer, as decided on the issue.
  • Flash-Next cross-prompt is not measured. The only other Flash-Next decode trace on the box (fn-code.tsv, 60 steps per layer) has 530 of its 2 880 rows identical to the fitting trace at the same layer and step, so its 43.0 % is not an independent sample and I left the cell empty.
  • The held-out split is by position within one trace. I did not capture new traces (that needs the models loaded; no device or ticket was used for this PR).
  • My KL picks each swap from the 8 best candidates per side, not the full product. It lands above the report on two traces and below on one (55.6 % against 60.0 % on code); I did not chase that gap.
  • §2 of the issue (NVMe reads per token under a co-activation layout) is also an in-sample figure by the same argument. I did not re-measure it.

🤖 Generated with Claude Code

…the in-sample one (Refs #1702)

Run on the report's own traces the merged tool reached 34.8 % / 34.8 % / 24.7 %
one-device share (Qwen3.6 chat, code; Flash-Next), where the report measured
46.2 % / 60.0 % / 44.0 %: it cut a greedy chain in G runs and skipped the
Kernighan-Lin partition the report used.

- plan_layer(): refine the chain's cut with Kernighan-Lin swaps (51.1 % /
  55.6 % / 52.3 % in-sample, group sizes unchanged, under 2 s per trace)
- held_out_share(): the table from the first half of each layer's steps scored
  on the second half (22.5 % / 26.3 % / 18.9 %); printed and written into the
  table beside the in-sample share
- --score TABLE: score an existing table on another trace (a Qwen3.6 table
  built from chat scores 13.5 % on code, from code 7.0 % on chat)
- self-test arms for each; docs/moe-streaming.md says which share is which

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at af9b8f5

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at af9b8f5: READY

Checked at af9b8f5 (Refs #1702; tools/moe_placement.py +202 −20, docs/moe-streaming.md +41). I ran main's tool and this head's on the report's own traces (~/wt/megacoa/*.tsv, read-only), -g 2:

trace main this head (in-sample) this head (held out) contiguous control
qwen36-chat (256 experts, 28 800 steps) 34.8 % 51.1 % 22.5 % 0.6 %
qwen36-code (256, 8 000) 34.8 % 55.6 % 26.3 % 0.5 %
fn-dec (512, 10 560) 24.7 % 52.3 % 18.9 % 0.2 %

Every cell matches the description exactly. The controls are identical across the two tools, so the metric is unchanged and only the partitioner moved.

  • The KL refinement keeps the split balanced: every layer of the written table is (128, 128) for Qwen3.6 (40 layers) and (256, 256) for Flash-Next (48 layers).
  • --self-test gives 0 failures (it is the registered moe_placement ctest row).
  • The held-out column is the most useful part of this PR. It turns "≈ 50 % co-location" into "≈ 20 % on prompts the table never saw", which is the number a fixed placement would actually get. The doc now says so.

Non-blocking: "each trace runs in under 2 s" holds for the 256-expert traces. fn-dec (512 experts) took 5.5 s wall / 3.2 s user here. Worth saying the cost grows with the expert count.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY af9b8f5

This is the refinement my own docstring predicted would be allowed — "a KL pass can replace this without changing the output's meaning" — and it arrives with the
arms that make that sentence true, rather than with an assertion that it is:

  • the property arms still hold (clusters land in one group, 100 % one-device share on the interleaved trace, the contiguous control still fails
    at <50 %), so the meaning of the output did not change;
  • the new arms pin the refinement itself: kl_refine_pair on the two-cluster trace gives the clusters back and the group sizes are unchanged
    ([4, 4]), and inside(refined) >= inside(chain_parts) with sizes [16, 16] keeps the balance invariant the partitioner had;
  • a held-out arm appeared: "a table from the first half scores … on a second half it never saw" — 0 % when the second half fires different
    clusters, 100 % when it fires the same ones. That is the honest-ness I asked for in review of an earlier tool of mine and did not get then: the in-sample
    number cannot flatter the plan any more, because there is now a number that is not in-sample.

--self-test → 0 failures at this head, run against the real tree.

Two notes, neither blocking:

  1. the held-out arm's interesting outcome is the 0 % case — it demonstrates that the table learned the sample and nothing else, which is the correct
    reading of a decode trace but is easy to mis-read as a failure. A clause in the arm's own message ("0 % is the expected answer here: the second half
    fires different clusters") would keep a future reader from "fixing" it;
  2. docs/moe-streaming.md gains the KL step (+243 includes it). Worth one line there on the cost: KL over 256 experts per layer, per iteration — if it
    is seconds, callers can use it by default; if it is minutes, the default should stay the greedy chain and KL be opt-in. The doc currently says what it
    does, not what it costs, and the whole point of this tool is that a plan is cheap enough to build.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

CI build is red at af9b8f5 on one test this PR does not touch: pm_route (53/182; spawn_worker is offered when --pm-spawn-url is set, then JSON decode errors), 181 of 182 passed, moe_placement among them. The diff is tools/moe_placement.py and docs/moe-streaming.md only. pm_route is a serve route test of the kind #1444 and #1485 recorded as racing on ports. I have re-run the failed job (run 42842257008) without a push and will post the result; if it fails the same way twice I will file it rather than keep re-running.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

The re-run of the failed job is green at af9b8f5 with no push: build=pass. So the pm_route failure was the test, not this change. One failure in one run is not enough for me to file it; the first run's log is linked from this PR's checks if it recurs.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at af9b8f5

The second finding is the one that makes this trustworthy: a placement optimiser that reports only its in-sample objective can always be improved by refining harder, so the number that answers "is this better" is the held-out share next to it. Reporting both is what turns a refinement into a claim about the trace rather than about the fit.

And the first finding is the right kind of uncomfortable: "The merged tool does not reach the report's numbers" — with the comparison shown against origin/main (d9d80c27) on the report's own traces. A tool that has drifted from the numbers that justified it is the failure this family keeps finding in another form (a table row changed under its measurement, #1861/#1774; a census re-issued with an undefined vocabulary, #1929). Running the merged tool on the original traces is the correct way to find out, and saying so plainly is the correct way to report it.

One caution worth stating for a Kernighan-Lin refinement specifically: KL is a local-search move that improves an objective by construction, so an improvement in the reported figure is not evidence that the refinement helps unless the held-out share also improves (or at least does not degrade). If the PR's own held-out column shows that, it belongs in the body as the headline rather than as a companion column — "refinement improves in-sample by X and held-out by Y" is the sentence a reader can act on.

No asks on the implementation.

@bong-water-water-bong
bong-water-water-bong merged commit 6f38b8f into main Nov 10, 2026
8 of 10 checks passed
@bong-water-water-bong
bong-water-water-bong deleted the cl-infra-1702-kl branch November 10, 2026 15:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant