From f0a8e727eb26ca5d362a63f2f768630af48e2f33 Mon Sep 17 00:00:00 2001 From: haoc <34830250+wimi321@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:34:27 -0700 Subject: [PATCH 1/2] Metal: hold the GPU batch size fixed to stop MPSGraph recompile stalls 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. --- cpp/configs/analysis_example.cfg | 7 +++++++ cpp/configs/gtp_example.cfg | 7 +++++++ cpp/neuralnet/metalbackend.cpp | 32 ++++++++++++++++++++++++++++---- cpp/neuralnet/metalbackend.h | 10 ++++++++++ 4 files changed, 52 insertions(+), 4 deletions(-) diff --git a/cpp/configs/analysis_example.cfg b/cpp/configs/analysis_example.cfg index cc09490414..a71a269683 100644 --- a/cpp/configs/analysis_example.cfg +++ b/cpp/configs/analysis_example.cfg @@ -260,6 +260,13 @@ nnRandomize = true # # Default (no config): 1 server thread, GPU-only mode (gpuIdx = 0). +# Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches. +# MPSGraph compiles a new executable, which takes seconds and stalls the search, the +# first time it sees each batch size. With padding (the default, auto) there is one +# shape, and the default nnMaxBatchSize becomes about half of numSearchThreads so the +# padding stays cheap. ANE (CoreML) threads never pad. +# metalPadBatch = auto + # FP16 precision (default true). Currently this only affects the ANE/CoreML path: # the GPU/MPSGraph path always runs in FP32 regardless of this flag. Setting this # to false also disables ANE acceleration (the ANE is FP16-only hardware, so diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 2d9610c7bc..8cbe3c0d21 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -585,6 +585,13 @@ searchFactorWhenWinningThreshold = 0.95 # # Default (no config): 1 server thread, GPU-only mode (gpuIdx = 0). +# Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches. +# MPSGraph compiles a new executable, which takes seconds and stalls the search, the +# first time it sees each batch size. With padding (the default, auto) there is one +# shape, and the default nnMaxBatchSize becomes about half of numSearchThreads so the +# padding stays cheap. ANE (CoreML) threads never pad. +# metalPadBatch = auto + # FP16 precision (default true). Currently this only affects the ANE/CoreML path: # the GPU/MPSGraph path always runs in FP32 regardless of this flag. Setting this # to false also disables ANE acceleration (the ANE is FP16-only hardware, so diff --git a/cpp/neuralnet/metalbackend.cpp b/cpp/neuralnet/metalbackend.cpp index e8b774fd71..e5bb2b1066 100644 --- a/cpp/neuralnet/metalbackend.cpp +++ b/cpp/neuralnet/metalbackend.cpp @@ -493,6 +493,7 @@ metalContext(createMetalComputeContext(nnX, nnY, useFP16Mode != enabled_t::False this->useFP16Mode = useFP16Mode; this->nnXLen = nnX; this->nnYLen = nnY; + this->padBatchMode = enabled_t::Auto; } ComputeContext::~ComputeContext() { @@ -524,10 +525,10 @@ ComputeContext* NeuralNet::createComputeContext( (void)logger; (void)homeDataDirOverride; (void)loadedModel; - (void)cfg; ComputeContext* context = new ComputeContext(nnXLen, nnYLen, useFP16Mode); context->aneOnly = aneOnly; + context->padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto; return context; } @@ -676,6 +677,7 @@ mpsGraphOnlyHandle(createMPSGraphHandleIfNeeded(context, loadedModel, requireExa coremlOnlyHandle(createCoreMLOnlyHandleIfNeeded(context, loadedModel, requireExactNNLen, maxBatchSize, gpuIdx, serverThreadIdx)) { bool hasMPSGraph = static_cast(mpsGraphOnlyHandle); bool hasCoreML = static_cast(coremlOnlyHandle); + padBatch = hasMPSGraph && context->padBatchMode != enabled_t::False; if(hasMPSGraph == hasCoreML) { throw runtime_error("Metal backend: Logic error - expected exactly one compute handle, got " + string(hasMPSGraph && hasCoreML ? "both" : "neither") + " (gpuIdx=" + to_string(gpuIdx) + ")"); } @@ -775,8 +777,12 @@ std::string NeuralNet::getRuntimeBackendDetail(ConfigParser& cfg) { } NeuralNet::BatchPolicy NeuralNet::getBatchPolicy(ConfigParser& cfg) { - (void)cfg; - return NeuralNet::BatchPolicy::Dynamic; + // MPSGraph compiles a new executable the first time it runs at each batch size, which takes + // seconds and stalls every search thread waiting on the GPU. So by default (metalPadBatch) the + // GPU path holds the batch size fixed and pads short batches, and the batch is sized the + // fixed-shape way (about half the concurrent evaluations) to keep the padding cheap. + enabled_t padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto; + return padBatchMode == enabled_t::False ? NeuralNet::BatchPolicy::Dynamic : NeuralNet::BatchPolicy::FixedShape; } int NeuralNet::getNumEffectiveDevices(ConfigParser& cfg, const std::vector& gpuIdxByServerThread) { @@ -1186,6 +1192,24 @@ void MetalProcess::getMetalOutput( batchSize); } else if(gpuHandle->mpsGraphOnlyHandle) { // GPU mode: Use MPSGraph (GPU) + // Run at a constant number of rows and throw the extra ones away, see getBatchPolicy. The + // padding rows copy row 0 rather than keeping whatever an earlier call left there (or zeros on + // the first call): an all-zero row has an empty mask, which the mask-weighted pooling divides by. + const int runBatchSize = gpuHandle->padBatch ? inputBuffers->maxBatchSize : batchSize; + for(int row = batchSize; row < runBatchSize; row++) { + copy( + inputBuffers->userInputBuffer, inputBuffers->userInputBuffer + inputBuffers->singleSpatialElts, + inputBuffers->userInputBuffer + inputBuffers->singleSpatialElts * row); + copy( + inputBuffers->userInputGlobalBuffer, inputBuffers->userInputGlobalBuffer + inputBuffers->singleInputGlobalElts, + inputBuffers->userInputGlobalBuffer + inputBuffers->singleInputGlobalElts * row); + copy( + inputBuffers->userInputMetaBuffer, inputBuffers->userInputMetaBuffer + inputBuffers->singleInputMetaElts, + inputBuffers->userInputMetaBuffer + inputBuffers->singleInputMetaElts * row); + copy( + inputBuffers->userInputMaskBuffer, inputBuffers->userInputMaskBuffer + inputBuffers->singleMaskElts, + inputBuffers->userInputMaskBuffer + inputBuffers->singleMaskElts * row); + } gpuHandle->mpsGraphOnlyHandle.get().apply( inputBuffers->userInputBuffer, inputBuffers->userInputGlobalBuffer, @@ -1195,7 +1219,7 @@ void MetalProcess::getMetalOutput( inputBuffers->valueResults, inputBuffers->scoreValuesResults, inputBuffers->ownershipResults, - batchSize); + runBatchSize); } else { throw runtime_error("Metal backend: No valid compute handle available"); } diff --git a/cpp/neuralnet/metalbackend.h b/cpp/neuralnet/metalbackend.h index e77dd18d9b..ad7d681303 100644 --- a/cpp/neuralnet/metalbackend.h +++ b/cpp/neuralnet/metalbackend.h @@ -113,6 +113,11 @@ struct ComputeContext { */ MetalComputeContext metalContext; + /** + * @brief metalPadBatch: whether GPU (MPSGraph) threads run every batch at the max batch size. + */ + enabled_t padBatchMode; + /** * @brief True only when every configured device is METAL_MUX_ANE, so no * MPSGraph (GPU) handle will ever read modelDesc weights. Gates the call to @@ -155,6 +160,11 @@ struct ComputeHandle { */ int gpuIndex; + /** + * @brief Whether this handle runs every MPSGraph batch at the max batch size (see metalPadBatch). + */ + bool padBatch; + /** * @brief The version of the loaded model. */ From 31faa715cce9941bd75e25b0a2ab2336ddcda5c5 Mon Sep 17 00:00:00 2001 From: haoc <34830250+wimi321@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:56:40 -0700 Subject: [PATCH 2/2] Metal: keep a configured nnMaxBatchSize dynamic 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. --- cpp/configs/analysis_example.cfg | 8 +++++--- cpp/configs/gtp_example.cfg | 7 ++++--- cpp/neuralnet/metalbackend.cpp | 22 ++++++++++++++++------ cpp/neuralnet/metalbackend.h | 2 +- 4 files changed, 26 insertions(+), 13 deletions(-) diff --git a/cpp/configs/analysis_example.cfg b/cpp/configs/analysis_example.cfg index a71a269683..4736c0afbb 100644 --- a/cpp/configs/analysis_example.cfg +++ b/cpp/configs/analysis_example.cfg @@ -262,9 +262,11 @@ nnRandomize = true # Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches. # MPSGraph compiles a new executable, which takes seconds and stalls the search, the -# first time it sees each batch size. With padding (the default, auto) there is one -# shape, and the default nnMaxBatchSize becomes about half of numSearchThreads so the -# padding stays cheap. ANE (CoreML) threads never pad. +# first time it sees each batch size. With padding there is one shape, and the default +# nnMaxBatchSize becomes about half of numAnalysisThreads * numSearchThreadsPerAnalysisThread +# so the padding stays cheap. The default (auto) pads only when nnMaxBatchSize is not set, +# since padding every batch up to a size chosen for dynamic batching wastes GPU work; this +# file sets nnMaxBatchSize above, so remove it to pad. ANE (CoreML) threads never pad. # metalPadBatch = auto # FP16 precision (default true). Currently this only affects the ANE/CoreML path: diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 8cbe3c0d21..ae006666c7 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -587,9 +587,10 @@ searchFactorWhenWinningThreshold = 0.95 # Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches. # MPSGraph compiles a new executable, which takes seconds and stalls the search, the -# first time it sees each batch size. With padding (the default, auto) there is one -# shape, and the default nnMaxBatchSize becomes about half of numSearchThreads so the -# padding stays cheap. ANE (CoreML) threads never pad. +# first time it sees each batch size. With padding there is one shape, and the default +# nnMaxBatchSize becomes about half of numSearchThreads so the padding stays cheap. +# The default (auto) pads only when nnMaxBatchSize is not set, since padding every batch +# up to a size chosen for dynamic batching wastes GPU work. ANE (CoreML) threads never pad. # metalPadBatch = auto # FP16 precision (default true). Currently this only affects the ANE/CoreML path: diff --git a/cpp/neuralnet/metalbackend.cpp b/cpp/neuralnet/metalbackend.cpp index e5bb2b1066..1d7981e1f8 100644 --- a/cpp/neuralnet/metalbackend.cpp +++ b/cpp/neuralnet/metalbackend.cpp @@ -493,12 +493,20 @@ metalContext(createMetalComputeContext(nnX, nnY, useFP16Mode != enabled_t::False this->useFP16Mode = useFP16Mode; this->nnXLen = nnX; this->nnYLen = nnY; - this->padBatchMode = enabled_t::Auto; + this->padBatch = true; } ComputeContext::~ComputeContext() { } +// metalPadBatch: whether GPU (MPSGraph) batches run padded up to nnMaxBatchSize, see getBatchPolicy. +static bool metalPadsBatch(ConfigParser& cfg) { + enabled_t padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto; + if(padBatchMode == enabled_t::Auto) + return !cfg.contains("nnMaxBatchSize"); + return padBatchMode == enabled_t::True; +} + ComputeContext* NeuralNet::createComputeContext( const vector& gpuIdxs, Logger* logger, @@ -528,7 +536,7 @@ ComputeContext* NeuralNet::createComputeContext( ComputeContext* context = new ComputeContext(nnXLen, nnYLen, useFP16Mode); context->aneOnly = aneOnly; - context->padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto; + context->padBatch = metalPadsBatch(cfg); return context; } @@ -677,7 +685,7 @@ mpsGraphOnlyHandle(createMPSGraphHandleIfNeeded(context, loadedModel, requireExa coremlOnlyHandle(createCoreMLOnlyHandleIfNeeded(context, loadedModel, requireExactNNLen, maxBatchSize, gpuIdx, serverThreadIdx)) { bool hasMPSGraph = static_cast(mpsGraphOnlyHandle); bool hasCoreML = static_cast(coremlOnlyHandle); - padBatch = hasMPSGraph && context->padBatchMode != enabled_t::False; + padBatch = hasMPSGraph && context->padBatch; if(hasMPSGraph == hasCoreML) { throw runtime_error("Metal backend: Logic error - expected exactly one compute handle, got " + string(hasMPSGraph && hasCoreML ? "both" : "neither") + " (gpuIdx=" + to_string(gpuIdx) + ")"); } @@ -780,9 +788,11 @@ NeuralNet::BatchPolicy NeuralNet::getBatchPolicy(ConfigParser& cfg) { // MPSGraph compiles a new executable the first time it runs at each batch size, which takes // seconds and stalls every search thread waiting on the GPU. So by default (metalPadBatch) the // GPU path holds the batch size fixed and pads short batches, and the batch is sized the - // fixed-shape way (about half the concurrent evaluations) to keep the padding cheap. - enabled_t padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto; - return padBatchMode == enabled_t::False ? NeuralNet::BatchPolicy::Dynamic : NeuralNet::BatchPolicy::FixedShape; + // fixed-shape way (about half the concurrent evaluations) to keep the padding cheap. An + // nnMaxBatchSize set in the config is normally sized for dynamic batching (analysis_example.cfg + // asks for at least the total number of search threads), and padding every batch up to it would + // multiply the GPU work, so by default such a config keeps dynamic batches. + return metalPadsBatch(cfg) ? NeuralNet::BatchPolicy::FixedShape : NeuralNet::BatchPolicy::Dynamic; } int NeuralNet::getNumEffectiveDevices(ConfigParser& cfg, const std::vector& gpuIdxByServerThread) { diff --git a/cpp/neuralnet/metalbackend.h b/cpp/neuralnet/metalbackend.h index ad7d681303..b4429a3c9a 100644 --- a/cpp/neuralnet/metalbackend.h +++ b/cpp/neuralnet/metalbackend.h @@ -116,7 +116,7 @@ struct ComputeContext { /** * @brief metalPadBatch: whether GPU (MPSGraph) threads run every batch at the max batch size. */ - enabled_t padBatchMode; + bool padBatch; /** * @brief True only when every configured device is METAL_MUX_ANE, so no