Skip to content

Metal: hold the GPU batch size fixed to stop MPSGraph recompile stalls - #1270

Open
wimi321 wants to merge 2 commits into
lightvector:masterfrom
wimi321:metal-fixed-batch-shape
Open

wimi321 wants to merge 2 commits into
lightvector:masterfrom
wimi321:metal-fixed-batch-shape

Conversation

@wimi321

@wimi321 wimi321 commented Oct 21, 2026 •

Copy link
Copy Markdown

MPSGraph compiles a new executable the first time graph.encode sees a new set of input shapes. The search varies the batch size from call to call, so every batch size not seen yet stalls all search threads while it compiles. That is about 3 s each for the b11 transformer (kata1-tf3-b11c768-s12002M) on an M4 Pro.

Restarting a search (a new kata-analyze, or moving to another position) makes the next batches small. So a GTP session hits several of these stalls in its first minutes, then never again. Running macOS sample during a stall shows:

  • the MPSGraphExecutable_queue thread in mlir::collectRegions;
  • every search thread waiting in NNEvaluator::evaluate.

When KataGo sizes the batch itself, this PR uses the existing FixedShape batch policy for Metal, the same way the ONNX backend does with onnxPadBatch:

  • With no nnMaxBatchSize in the config, getBatchPolicy returns FixedShape. The default nnMaxBatchSize becomes about half the concurrent evals (computeFixedShapeMaxBatchSize), which the Metal section of the example configs already recommends.
  • GPU (MPSGraph) handles then run every batch at nnMaxBatchSize and discard the extra rows. Padding rows copy row 0, since an all-zero row has an empty mask.
  • A configured nnMaxBatchSize keeps Dynamic batches. It is usually sized for dynamic batching (analysis_example.cfg sets 64), and padding every batch up to it multiplies the work (second table below).
  • ANE (CoreML) handles evaluate row by row and never pad.
  • metalPadBatch = true pads a configured size too, and metalPadBatch = false restores the previous behavior. Both are documented in gtp_example.cfg and analysis_example.cfg.

GTP measurement setup:

  • one start, then 24 stop + kata-analyze restarts on one midgame position, alternating with and without focus;
  • b11 network, M4 Pro;
  • a minimal GTP config: numSearchThreads as listed, and every batch setting left at its default.
stalls after startup first visits after start visits/s
master, 8 threads 5 6.1 s 57
this PR, 8 threads 0 3.1 s 54
master, 6 threads 3 6.0 s 55
this PR, 6 threads 0 3.1 s 55
this PR, metalPadBatch = false 5 6.0 s 58

Padding to the Dynamic default size instead (8 rows) halved throughput (29 visits/s), which is why the batch is sized the FixedShape way.

Analysis engine: analysis_example.cfg (nnMaxBatchSize = 64) with numSearchThreadsPerAnalysisThread = 1, the first 80 positions of one game sent as separate 2-visit queries, kata1-tf3-b11c768-s12002M, M4 Pro:

80 queries first result
master 6.8 s 2.9 s
first commit of this PR (padded to the configured 64) 212.1 s 6.7 s
this PR 6.8 s 2.9 s

With the second commit, the 6-thread GTP run still has no stalls (first visits 3.1 s, 54 visits/s), and with nnMaxBatchSize = 16 set it behaves like master (3 stalls).

testbackendreference (b11) passes, with error statistics matching master to within 0.00003, and runnnlayertests passes. I measured on arm64 only. The change also builds on an Intel (x86_64) macOS runner, where we ship it in our own engine build, but I had no Intel Mac to measure on.

MPSGraph compiles a new executable the first time it runs at each input
shape, and the search varies the batch size from call to call, so every new
batch size stalled all search threads for seconds (about 3s each for a b11
net on an M4 Pro). Restarting a search, e.g. a new kata-analyze, makes the
next batches small, so this hit repeatedly early in each session.

Use the existing FixedShape batch policy for the Metal backend, as the ONNX
backend does with onnxPadBatch: GPU (MPSGraph) threads run every batch at
nnMaxBatchSize, filling the padding rows with copies of row 0, and that size
now defaults to about half the concurrent evals so the padding stays cheap.
ANE (CoreML) threads evaluate row by row and never pad. metalPadBatch = false
restores the previous behavior.
Padding every GPU batch up to nnMaxBatchSize multiplies the work when the config
sizes it for dynamic batching. analysis_example.cfg sets 64, so an analysis
engine with two search threads (numAnalysisThreads = 2,
numSearchThreadsPerAnalysisThread = 1, 2 visits per query) ran 64 rows for every
one or two positions: 80 queries took 217 s instead of 7 s on an M4 Pro.

By default, pad only when KataGo derives the batch size, i.e. the config does not
set nnMaxBatchSize. metalPadBatch = true still pads a configured size.

This branch has not been deployed

No deployments
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