mlxrunner: raise the MTP pending-flush cap to 256 tokens

Per-token cost of the batched head forward keeps falling until the
flush is large enough to reach the fastest kernels: NAX matmul tiles
for dense heads, and the segmented gather path for MoE heads, which
needs tokens*topK/experts >= 4. Measured across the qwen3.6 heads,
256 is the smallest cap past every threshold and within a few percent
of each head's per-token floor. The cost is bounded: up to 2.5 MiB of
pinned hiddens per request and a flush stall under one decode step.
This commit is contained in:
Jesse Gross
2026-07-07 12:50:19 -07:00
parent 556245843a
commit 123b1f2479
+1 -1
View File
@@ -11,7 +11,7 @@ import (
// mtpPendingFlushTokens caps how many committed look-ahead tokens wait in the
// pending buffer before a batched flush, bounding the pinned hidden states
// regardless of what else triggers a flush.
const mtpPendingFlushTokens = 32
const mtpPendingFlushTokens = 256
// mtpDrafter drafts with a model's multi-token-prediction head. Constructed
// at load, it fixes the trie keys' draft look-ahead for the model's lifetime