From e3fb72adbc546bad07bea137b705f109c90a0dcc Mon Sep 17 00:00:00 2001 From: Aegis-AI Date: Tue, 12 May 2026 15:23:38 -0700 Subject: [PATCH 1/2] Fix GPU Hang: flush Metal graph with eval() after MTPTokenIterator verification pass On hybrid SSM/attention models (Qwen35), the recurrent GatedDeltaNet layers accumulate un-evaluated MLX graph nodes across each speculateRound(). Without an explicit eval() after callMTP(), the Metal command buffer grows across multiple speculation rounds until it triggers the GPU Watchdog (kIOGPUCommandBufferCallbackErrorHang). Adding eval(mtpResult) immediately after the verification forward pass flushes the accumulated graph, preventing the Metal timeout. --- Libraries/MLXLMCommon/Evaluate.swift | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index 07d73dd06..11c7e1e7b 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -1227,7 +1227,13 @@ public struct MTPTokenIterator: TokenIteratorProtocol { let mtpResult = model.callMTP(verifyInput.tokens[.newAxis], cache: cache, mtpCaches: mtpCaches) guard !mtpResult.isEmpty else { return } - + + // Flush the Metal command buffer immediately after the verification forward pass. + // On hybrid SSM/attention models (e.g. Qwen35), the recurrent SSM layers accumulate + // un-evaluated graph nodes across rounds. Without an explicit sync here the Metal + // command buffer grows until it triggers the GPU Watchdog. + eval(mtpResult) + let mainLogits = mtpResult[0] let mainTokens: MLXArray From 881cddfc3d7b57b69140492151512966ee12c1c6 Mon Sep 17 00:00:00 2001 From: Simba Date: Tue, 12 May 2026 15:48:35 -0700 Subject: [PATCH 2/2] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- Libraries/MLXLMCommon/Evaluate.swift | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/Libraries/MLXLMCommon/Evaluate.swift b/Libraries/MLXLMCommon/Evaluate.swift index 11c7e1e7b..e816594a4 100644 --- a/Libraries/MLXLMCommon/Evaluate.swift +++ b/Libraries/MLXLMCommon/Evaluate.swift @@ -1228,13 +1228,16 @@ public struct MTPTokenIterator: TokenIteratorProtocol { let mtpResult = model.callMTP(verifyInput.tokens[.newAxis], cache: cache, mtpCaches: mtpCaches) guard !mtpResult.isEmpty else { return } + let mainLogits = mtpResult[0] + // Flush the Metal command buffer immediately after the verification forward pass. // On hybrid SSM/attention models (e.g. Qwen35), the recurrent SSM layers accumulate // un-evaluated graph nodes across rounds. Without an explicit sync here the Metal // command buffer grows until it triggers the GPU Watchdog. - eval(mtpResult) - - let mainLogits = mtpResult[0] + // + // Only force the main logits needed for verification/sampling so we avoid eagerly + // evaluating speculative MTP head logits that may be discarded on rejection. + eval(mainLogits) let mainTokens: MLXArray var mainProcessedLogits = [MLXArray]()