feat(qwen3-next): add pipeline-parallel stage support - #580
Merged
Conversation
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.
16 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds pipeline-parallel (PP) stage support for
qwen3_next, which was the only Qwen3-Next-style hybrid family missing it (qwen3_5already has full PP support and reusesqwen3_next's building blocks).Changes
Qwen3NextStageModelinsrc/models/qwen3_next.rs, modeled onQwen35StageModel: loads only the layers infilter.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.Qwen3NextStageExecutorinsrc/distributed/pipeline/stage_executor/qwen3_next.rs, modeled onqwen35.rs.src/distributed/pipeline/stage_executor/mod.rs: theQwen3NextStageFamilyvariant, name mapping, PP-capable set, andModelTypedispatch.Correctness
qwen3_nextmixes 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.qwen3_nextcheckpoint. Noqwen3_nextcheckpoint 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