Repository navigation
Conversation
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.
7 of 13 tasks
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MPSGraph compiles a new executable the first time
graph.encodesees 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 macOSsampleduring a stall shows:MPSGraphExecutable_queuethread inmlir::collectRegions;NNEvaluator::evaluate.When KataGo sizes the batch itself, this PR uses the existing
FixedShapebatch policy for Metal, the same way the ONNX backend does withonnxPadBatch:nnMaxBatchSizein the config,getBatchPolicyreturnsFixedShape. The defaultnnMaxBatchSizebecomes about half the concurrent evals (computeFixedShapeMaxBatchSize), which the Metal section of the example configs already recommends.nnMaxBatchSizeand discard the extra rows. Padding rows copy row 0, since an all-zero row has an empty mask.nnMaxBatchSizekeepsDynamicbatches. It is usually sized for dynamic batching (analysis_example.cfgsets 64), and padding every batch up to it multiplies the work (second table below).metalPadBatch = truepads a configured size too, andmetalPadBatch = falserestores the previous behavior. Both are documented ingtp_example.cfgandanalysis_example.cfg.GTP measurement setup:
stop+kata-analyzerestarts on one midgame position, alternating with and withoutfocus;numSearchThreadsas listed, and every batch setting left at its default.metalPadBatch = falsePadding to the
Dynamicdefault size instead (8 rows) halved throughput (29 visits/s), which is why the batch is sized theFixedShapeway.Analysis engine:
analysis_example.cfg(nnMaxBatchSize = 64) withnumSearchThreadsPerAnalysisThread = 1, the first 80 positions of one game sent as separate 2-visit queries,kata1-tf3-b11c768-s12002M, M4 Pro:With the second commit, the 6-thread GTP run still has no stalls (first visits 3.1 s, 54 visits/s), and with
nnMaxBatchSize = 16set it behaves like master (3 stalls).testbackendreference(b11) passes, with error statistics matching master to within 0.00003, andrunnnlayertestspasses. 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.