Skip to content

feat(qwen3-next): add pipeline-parallel stage support - #580

Merged
inureyes merged 1 commit into
mainfrom
feature/issue-510-qwen3-next-pp-stage
Jul 1, 2026
Merged

feat(qwen3-next): add pipeline-parallel stage support#580
inureyes merged 1 commit into
mainfrom
feature/issue-510-qwen3-next-pp-stage

Conversation

@inureyes

@inureyes inureyes commented Jul 1, 2026

Copy link
Copy Markdown
Member

Summary

Adds pipeline-parallel (PP) stage support for qwen3_next, which was the only Qwen3-Next-style hybrid family missing it (qwen3_5 already has full PP support and reuses qwen3_next's building blocks).

Changes

  • Qwen3NextStageModel in src/models/qwen3_next.rs, modeled on Qwen35StageModel: loads only the layers in filter.layer_range, accepts/produces hidden states on intermediate stages and token IDs / final logits on the first/last stage, and ties embeddings on the final stage.
  • Qwen3NextStageExecutor in src/distributed/pipeline/stage_executor/qwen3_next.rs, modeled on qwen35.rs.
  • Registration in src/distributed/pipeline/stage_executor/mod.rs: the Qwen3Next StageFamily variant, name mapping, PP-capable set, and ModelType dispatch.

Correctness

qwen3_next mixes linear (gated-delta, ArraysCache) and attention (KVCache) layers. The stage model builds and advances the correct cache type per local layer and derives the SSM/attention masks from the layers actually present in the stage, using global layer indices for typing. Covered by new unit tests: stage layer typing on global indices, per-stage attention offset (first-stage local attention and the all-linear default), and cache snapshot/restore round-trips for both the linear and attention caches.

Validation

  • cargo check --lib --tests --features metal,accelerate: compiles clean.
  • cargo test --release --features metal,accelerate qwen3_next: 10 tests pass (helper layout, stage typing/offset, cache round-trips, plan-generator replication, kv-arch interval, family registration).
  • cargo test --release --features metal,accelerate family_registry: 5 pass.
  • Not validated here: a real multi-stage run on an actual qwen3_next checkpoint. No qwen3_next checkpoint is present in this environment (models/ is unpopulated), so the end-to-end multi-stage decode parity against single-process decode was not executed and should be confirmed on hardware with a real checkpoint before relying on it in production.

Closes #510

Add Qwen3NextStageModel and Qwen3NextStageExecutor so qwen3_next checkpoints shard across pipeline-parallel stages, mirroring the existing qwen3_5 stage plumbing. Each stage loads only its layer range, builds and advances the correct per-layer cache type (gated-delta ArraysCache for linear layers, KVCache for attention layers), and derives the SSM and attention masks from the stage-local layers using global layer indices. Registers the Qwen3Next StageFamily variant with name mapping, PP-capable set, and ModelType dispatch.
@inureyes inureyes added the status:done Completed label Jul 1, 2026
@inureyes
inureyes merged commit 8efd020 into main Jul 1, 2026
7 checks passed
@inureyes
inureyes deleted the feature/issue-510-qwen3-next-pp-stage branch July 1, 2026 11:37
@inureyes inureyes self-assigned this Jul 1, 2026
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions priority:medium Medium priority labels Sep 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

priority:medium Medium priority status:done Completed type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant