feat(qwen3-next): add pipeline-parallel stage support (StageModel + stage executor) #510

Summary

mlxcel supports pipeline-parallel (PP) layer-range sharding across stages for many families, but qwen3_next is missing it. The closely related qwen3_5 family already has full PP stage support, so qwen3_next is the odd one out among the Qwen3-Next-style hybrid (gated-delta + attention) models.

Current state in mlxcel

• PP infrastructure exists: StageFamily enum and per-model dispatch in src/distributed/pipeline/stage_executor/mod.rs (families listed at mod.rs:241-244, ModelType -> StageFamily match at mod.rs:331-354), with per-family executors under src/distributed/pipeline/stage_executor/ (e.g. qwen35.rs, deepseek_v3.rs, gemma4.rs, gpt_oss.rs, llama.rs, ...).
• qwen3_5 is fully PP-shardable: Qwen35StageModel (src/models/qwen3_5.rs:2079-2156) with stage_index / filter.layer_range / hidden-state and token-ID stage I/O, plus the dedicated executor src/distributed/pipeline/stage_executor/qwen35.rs.
• qwen3_next has no stage support: src/models/qwen3_next.rs implements no StageModel/LayerFilter/StageExecutionOutput, and there is no stage_executor/qwen3_next.rs. So a qwen3_next checkpoint cannot be served across PP stages.

This matters because qwen3_5 itself reuses qwen3_next's building blocks (Qwen3NextAttention, Qwen3NextMLP, Qwen3NextRMSNormGated, gated_delta_update), so the PP work for qwen3_next is structurally very close to the existing qwen3_5 stage implementation.

What needs to be built

1. Qwen3NextStageModel in src/models/qwen3_next.rs, mirroring Qwen35StageModel: load only the layers in filter.layer_range, accept/produce hidden states (intermediate stages) or token IDs / final logits (first/last stage), tie embeddings on the final stage.
2. Hybrid cache across stage boundaries. qwen3_next mixes linear (gated-delta, ArraysCache) and attention (KVCache) layers; the stage model must build and advance the correct cache type per local layer, exactly as the single-process path does, and the SSM/attention masks must be derived from the layers actually present in the stage.
3. Qwen3NextStageExecutor under src/distributed/pipeline/stage_executor/qwen3_next.rs, modeled on qwen35.rs.
4. Registration: add the StageFamily variant(s) in stage_executor/mod.rs (enum, name mapping, PP-capable set, and the ModelType::Qwen3Next* -> StageFamily match arms), covering the MoE/VLM variants if applicable.

Acceptance criteria

• qwen3_next shards across >= 2 PP stages and completes generation.
• Multi-stage decode output matches single-process decode for the same checkpoint/prompt/seed (jitter-aware parity, consistent with how other hybrid families are gated).
• Hybrid cache (gated-delta ArraysCache + attention KVCache) is constructed and advanced correctly per stage; masks derived from stage-local layers.
• qwen3_next appears in the PP-capable family set and is dispatched to the new executor.
• Validated on a real qwen3_next checkpoint across stages.

Files

• src/models/qwen3_next.rs (new Qwen3NextStageModel)
• src/distributed/pipeline/stage_executor/qwen3_next.rs (new executor)
• src/distributed/pipeline/stage_executor/mod.rs (StageFamily variant + dispatch + PP-capable set)
• tests mirroring the qwen3_5 stage tests

Context

Discovered while auditing pipeline-parallel coverage during an upstream-sync review: qwen3_5 gained model-level pipelining but qwen3_next was never given the equivalent stage plumbing in mlxcel.