From 123b1f24793f4d0edc77e09dfcafb4c23b94ea72 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Tue, 7 Jul 2026 12:50:19 -0700 Subject: [PATCH] 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. --- x/mlxrunner/mtp.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x/mlxrunner/mtp.go b/x/mlxrunner/mtp.go index d559c6b4..bf40e569 100644 --- a/x/mlxrunner/mtp.go +++ b/x/mlxrunner/mtp.go @@ -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