Skip to content

v0.2.6 — a K3 MoE block computes from the arena

Choose a tag to compare

@pjordanandrsn pjordanandrsn released this 30 Jul 05:16
v0.2.6

0.2.5 gave the NVMe arena a compute consumer. Nothing called it — moe_layer_forward's only callers were its own tests. This is the model-side link.

The patch point

KimiSparseMoeBlock.moe_infer turned out to be a better fit than expected: upstream already builds exactly what the grouped kernel wants.

idxs = topk_ids.view(-1).argsort()
sorted_tokens = x[idxs // topk_ids.shape[1]]   # group-sorted [T, K]  == a_cat
tokens_per_expert = cnts.sum(dim=0)            #                      == sizes
for i, num_tokens in enumerate(tokens_per_expert):
    expert_out = expert(tokens_for_this_expert)   # one matmul PER expert

That per-expert loop — each iteration first materialising an expert's weights — is precisely what gemm_mxfp4_grouped collapses into one launch on packed bytes.

So the patch swaps the loop and nothing else. Sorting, weighting, unsorting and the shared-expert path remain upstream's; patching those would change the model rather than accelerate it. use_latent_moe's projections and the shared experts live in forward(), outside moe_infer, and are untouched.

Honest instrumentation

arena_call_stats() reports patched and calls separately, because a patch count is not a call count — both inert fixes in this project's history patched cleanly and never executed. patched=3, calls=0 is an expressible state, and a gate pins that zero before a forward, with a positive control that moves it.

Layer 0 (dense first_k_dense_replace) is skipped. Layers missing from a partial bake stay on resident experts rather than being silently routed to nothing. Patching nothing raises rather than returning 0.

Still required for K3 inference

A baked arena (1446.5 GB relocation — measured, not estimated) and the always-active weights resident. This release makes the model able to use an arena; it does not create one.