GPU watchdog kills process during long-context SDPA prefill (65K+ keys)

[Thump604]

Summary

steel_attention (full attention, qL > 8) dispatches all keys in a single Metal compute command. At 65K+ key sequence length, this dispatch exceeds the macOS GPU watchdog threshold (~5 seconds), resulting in kIOGPUCommandBufferCallbackErrorImpactingInteractivity and process termination.

This is distinct from #3267 (display-active watchdog during training with short sequences). This issue is about the absolute GPU time of a single SDPA dispatch during inference prefill at long context.

Reproduction

Any model with long context prefill triggers this. Tested on M2 Ultra 128GB, macOS 26.3.1:

• Qwen3.5-2B at 262K: GPU watchdog kill during prefill
• Qwen3.5-35B at 262K with YaRN: GPU watchdog kill during prefill
• Qwen3.5-122B at 128K+: borderline (single dispatch ~4-5s, close to threshold)

The vector SDPA path (qL ≤ 8, decode) is not affected — it already has 2-pass chunking.

Root Cause

sdpa_full_self_attention in scaled_dot_product_attention.cpp dispatches a single steel_attention kernel covering the entire key sequence. The dispatch grid is proportional to kL (number of key tokens). At 65K+ keys with typical head counts (8-64 GQA heads), the single dispatch runs for >5 seconds on Apple Silicon, triggering the macOS GPU watchdog.

The watchdog fires per-command-buffer, not per-dispatch. Multiple dispatches within the same command buffer don't individually reset the timer.

What I Tried

1. Mid-primitive command buffer commits: Inserting mx.eval() or command buffer commits between SDPA chunks. Triggers Metal assertion: "A command encoder is already encoding to this command buffer". MLX's eval model doesn't support mid-graph command buffer boundaries.

2. AGX_RELAX_CDM_CTXSTORE_TIMEOUT=1: Works as a workaround (per [BUG] Metal GPU watchdog kills LoRA training when display is active #3267), but it's a system-wide env var that disables the watchdog entirely — not appropriate as a permanent solution.

3. Adaptive prefill step size (application-level): Reducing prefill_step_size in the LLM framework so each generate() call processes fewer tokens. This creates separate mx.eval() calls and thus separate command buffers. Partially works, but adds significant overhead and doesn't help when a single SDPA layer's key length exceeds 65K (which happens during later prefill steps as KV cache accumulates).

Proposed Solution: Chunked Full Attention

Split the key sequence into chunks (e.g., 65K tokens each), dispatch steel_attention per chunk with partial outputs, then reduce via online softmax. This mirrors the existing sdpa_vector_2pass pattern already used for decode.

Design:

• Each chunk dispatches steel_attention with write_partial=true → outputs unnormalized O + partial max/sum
• Reduction merges chunk results: O_final = Σ(O_chunk × exp2(max_chunk - max_global))
• Routes automatically when kL ≥ threshold
• Single-pass path unchanged for short sequences

I have a working implementation on a local branch (5 commits, ~575 lines across 3 files). Verified at 512K context with Qwen3.5-2B. Happy to open a PR if this approach aligns with the project's direction.

Key design question: Should this be part of the existing steel_attention infrastructure, or a separate kernel? The existing sdpa_vector_2pass sets a precedent for chunking within the SDPA dispatch.

Related

• [BUG] Metal GPU watchdog kills LoRA training when display is active #3267 — GPU watchdog during training (display-active, different trigger)
• fix: add head_dim=256 to fused SDPA full attention kernel #3293 — head_dim=256 fused SDPA (fixes crash at 32K, but 65K+ still hits watchdog)
• SIGSEGV in QuantizedMatmul::eval_gpu during long token generation on Mac Studio M2 Ultra #3216 — thread-safety crashes (fixed by Make each thread have its own default stream #3281)

[hnshah]

Validation on M3 Ultra 256GB

I've independently validated your feat/chunked-sdpa fix on different hardware:

Test Hardware:

• Mac Studio M3 Ultra (256GB)
• macOS 25.3.0 (Darwin 25.3.0)
• MLX: 0.31.2.dev20260324+22525b51 (your branch)

Test Results:

Context Length	Head Dim	Time	Result
65K tokens	D=128	0.03s	✅ No watchdog kill
128K tokens	D=128	0.04s	✅ No watchdog kill
128K tokens	D=256	0.16s	✅ No watchdog kill
262K tokens	D=128	0.05s	✅ No watchdog kill

Environment:

MLX_SDPA_CHUNK_THRESHOLD=65536
MLX_SDPA_CHUNK_SIZE=65536

Validation:

• ✅ All contexts tested successfully (65K, 128K, 262K)
• ✅ No GPU watchdog kills at any tested length
• ✅ Output correctness verified (finite values, expected shapes)
• ✅ D=256 head dimension tested (Qwen3.5 use case)
• ✅ Performance is acceptable (< 0.2s for all tests)

Conclusion:

Your chunked SDPA implementation successfully solves the GPU watchdog issue on M3 Ultra. The fix works as described across all tested context lengths up to 262K tokens.

Recommendation: This is ready for PR submission. The implementation is solid and solves a critical blocker for long-context inference on Apple Silicon.

---

Technical notes:

• Tested with chunking enabled via environment variables
• All tests used float16 dtype (production use case)
• M3 Ultra provides independent validation beyond your M2 Ultra testing
• 256GB memory allowed testing up to 262K without OOM

Thank you for implementing this fix! Would you like help with the PR submission or review process?

[hnshah]

@Thump604 Your implementation is ready for PR! I've validated it works perfectly on M3 Ultra across all tested context lengths (65K-262K). The chunked SDPA approach is solid and solves the GPU watchdog issue completely.

If you'd like, I'm happy to help with:

• Code review
• Additional testing on other hardware configurations
• Documentation/comments refinement

Let me know if there's anything I can do to help get this merged! This is a critical fix for long-context users.

[Thump604]

Thanks @hnshah, really appreciate the independent validation on M3 Ultra, especially 262K - that's beyond what we can test on M2 Ultra 128GB.

The PRs are up:

• Add logsumexp output to fused SDPA kernel #3306 logsumexp output (prerequisite)
• Chunked full-attention SDPA for long key sequences #3307 chunked SDPA (what you validated)
• fix: add head_dim=256 to fused SDPA full attention kernel #3293 head_dim=256 routing (updated today per @angeloskath's feedback)

Would be great if you could post your results on #3306 and #3307 so the maintainers see them during review. And if you have Qwen3.5 models (head_dim=256), testing #3293 would help too since @jagrit06 was tagged to finalize the routing threshold.