Layr-Labs / Layr-Labs/d-inference
Provider MLX continuous-batching crash: [concatenate] shape mismatch merging concurrent requests (Qwen3.5 hybrid model)
- Dominant language
- Go
- Stars
- 522
- Forks
- 113
- Avg merge
- 17h 26m
- Merged PRs (30d)
- 111
Description
## Summary
Under concurrent requests to a single model, the Swift provider's MLX continuous-batching engine crashes with a fatal `[concatenate]` shape-mismatch, killing the provider process. All in-flight requests then fail with `502 provider_error: "inference failed after N attempt(s): provider disconnected"`. It is **intermittent** — `master` passes the same test, but it recurs across PR branches.
## The crash
```
MLX/ErrorHandler.swift:345: Fatal error: [concatenate] All the input array dimensions
must match exactly except for the concatenation axis. However, the provided shapes are
(2,8,24,64), (1,8,24,192), and the concatenation axis is -1.
at libs/mlx-swift/Source/Cmlx/mlx-c/mlx/c/ops.cpp:677
```
The two operands are `[batch, kv_heads, seq, head_dim]`:
- `(2, 8, 24, 64)` — batch 2, head_dim **64**
- `(1, 8, 24, 192)` — batch 1, head_dim **192**
Concatenating on axis `-1` requires every *other* axis to match, but the **batch dims differ (2 vs 1)** *and* the head_dim differs (64 vs 192). Two KV-cache tensors of incompatible geometry are being merged in the same `concatenated` call.
## Why this is a real model-shape collision (not just a transient)
`mlx-community/Qwen3.5-0.8B-MLX-4bit` is a **hybrid architecture**. The provider's own KV estimator classifies its layers into distinct geometries (`BatchScheduler+KVEstimation.swift`):
- `linear_attention` → GatedDeltaNet **recurrent-state** layers (`recurrentLayerTypes`, ~line 39-41)
- full-attention layers with `head_dim` vs `global_head_dim` and a `sliding_window_pattern` (lines 86-99)
So `head_dim 64` vs `192` corresponds to genuinely different per-layer cache/state shapes. When two concurrent requests at **different stages** (one mid-prefill, one mid-decode — hence batch 1 vs 2) are merged into one `GenerationBatch`, the per-layer cache merge concatenates tensors whose item-shapes don't line up.
## Suspected location (hypothesis — needs confirmation)
The batch-merge path lives in the **`libs/mlx-swift-lm` submodule** (Layr-Labs fork), not in `d-inference` directly:
- `Libraries/MLXLMCommon/ContinuousBatching/Scheduler.swift` — `mergeIntoGenBatch(_:)` merges a newly-prefilled generation into the running batch.
- `Libraries/MLXLMCommon/GenerationBatch.swift` — `extend(_ other:)` per-layer merge.
- `Libraries/MLXLMCommon/KVCache.swift` — `ArraysCache.extend(other:)` → `concatenateOptional(...)` builds zero-padded operands from `itemShape = shape.dropFirst()`, taking the shape from whichever operand is non-nil. If the two operands have different item-shapes (different head_dim / recurrent-state dims across layer types), the resulting `MLX.concatenated([lhs, rhs])` mismatches.
> Note: `concatenateOptional` calls `MLX.concatenated([...])` with **no axis** (defaults to axis 0), while the crash reports **axis -1** — so the exact failing call may be a *different* concat in the attention/cache-update path rather than this one. The mechanism (incompatible per-layer geometries merged across requests at different prefill/decode stages) is the same regardless of which concat trips first. The submodule code should be the focus of the fix.
There appears to be **no guard** that two requests are shape-compatible (same per-layer cache geometry, same stage) before they're merged into a shared batched forward pass.
## Reproduction
- Test: `TestProfile_SingleProviderNonStreaming` (`e2e/profile_test.go`), which fires **3 concurrent non-streaming requests** at a single Qwen3.5 provider.
- Observed: provider emits the fatal error mid-batch, disconnects; `errors (10/10) … status 502 … provider disconnected`; the test's `require.Greater(SuccessCount, 0)` trips with `"0" is not greater than "0"`.
- Intermittent: timing-dependent on whether request N's prefill completes and merges while requests 1..N-1 are mid-decode.
## Evidence it's flaky, not a branch regression
- **`master` ran the identical test and passed** (run `27508483291`, `TestProfile_SingleProviderNonStreaming` PASS in 22.1s).
- Recurs on multiple unrelated PR branches (`fix/coordinator-plaintext-body-cap`, `feat/oom-protection-vlm-media-cap`, `worktree-provider-trust-reliability`, `devin/*`) whose diffs don't touch the inference/KV path.
- Confirmed crash line on run `27510260770` / job `81308924460`.
## Impact
- **Provider availability**: a single crash takes down the provider process and 502s every co-tenant in-flight request — exactly the multi-tenant blast radius the coordinator's pre-content failover and cancellation logic try to contain. A retry lands on the same (now-dead) provider until the registry evicts it.
- **CI flakiness**: red E2E runs on PRs that didn't touch inference, eroding signal.
## Suggested fix directions
1. **Guard the batch merge**: before merging a generation into a running `GenerationBatch`, assert per-layer cache item-shapes are identical; if not, keep it in a separate batch instead of concatenating. (Same model → same architecture → shapes *should* match, so a mismatch indicates a real merge-ordering/stage bug to fix, not just to skip.)
2. **Fix `concatenateOptional` item-shape derivation** so it never builds a zero operand whose item-shape disagrees with the real operand (don't take `itemShape` from "whichever is non-nil"); and add a precondition that both operands' non-concat axes match.
3. **Fail soft, not fatal**: a per-request inference error should reject *that* request (or split the batch), not `fatalError` the whole provider process. A guard that throws a Swift error here would convert a fleet-impacting crash into a single 4xx/5xx.
## Notes
- The crashing concat is in the **`libs/mlx-swift-lm` submodule** — a fix likely lands there and is then pulled into `d-inference` via a submodule bump.
- Filing this so the batching crash gets a real fix rather than being masked by CI retries; not a blocker for coordinator-only PRs whose other E2E tests pass.
Contributor guide
Research direction
Start in the mlx-swift-lm submodule at Libraries/MLXLMCommon/ContinuousBatching/Scheduler.swift, then trace GenerationBatch.swift and KVCache.swift through mergeIntoGenBatch, extend, and concatenateOptional. Run TestProfile_SingleProviderNonStreaming in e2e/profile_test.go while inspecting the failing concurrent batch path. Done means incompatible cache geometries no longer reach a fatal concatenate, and concurrent requests do not disconnect the provider.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- swift
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 43/100