From f1433c552baacc590b1389b50ab86d6cc1ff7c08 Mon Sep 17 00:00:00 2001 From: Troll863 Date: Wed, 24 Jun 2026 21:08:16 +0800 Subject: [PATCH 1/4] docs(qwen35): draft TP2 phased design --- docs/index.md | 1 + docs/models/qwen35/tp-design.md | 418 ++++++++++++++++++++++++++++++++ 2 files changed, 419 insertions(+) create mode 100644 docs/models/qwen35/tp-design.md diff --git a/docs/index.md b/docs/index.md index 8cce9e77..9c7168ce 100644 --- a/docs/index.md +++ b/docs/index.md @@ -47,6 +47,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen35/model-crate.md` | `openinfer-qwen35-4b` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35-4b` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | | `models/qwen35/kernel-plan.md` | Qwen3.5-4B has a `openinfer_qwen35_4b::kernel_plan()` static descriptor mirroring the qwen3 module — enumerates every prefill/decode/unified op with its Rust call site, backend, and notes, so you can dump the active kernel mix without reading call sites. Pure refactor (issue #256), no kernel behavior change. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | +| `models/qwen35/tp-design.md` | Qwen3.5 TP2 design: reuse Qwen3's controller/worker TP runtime; Phase 1 shards full-attn+MLP while replicating linear/GDR, Phase 2 shards linear attention and GDR state using vLLM's Qwen3Next/GDN contract as reference. | ## models / deepseek-v4 diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md new file mode 100644 index 00000000..60c4b92c --- /dev/null +++ b/docs/models/qwen35/tp-design.md @@ -0,0 +1,418 @@ +# Qwen3.5 Tensor Parallelism Design + +> **TL;DR:** Qwen3.5 TP2 should reuse Qwen3's controller/worker TP runtime. Phase 1 shards full-attention and MLP while keeping linear attention/GDR replicated, so the dense TP path can be validated first; Phase 2 then shards linear attention, conv state, GDR recurrent state, and the GDR kernels using vLLM's Qwen3Next/GDN TP contract as the main reference. +> +> **Last touched:** 2026-06 + +## Goal + +Add single-node `TP=2` support for `Qwen3.5-4B`. + +The first goal is correctness and architectural integration, not peak performance. + +Concrete goals: + +- keep `TP=1` healthy +- add a fail-closed `TP=2` path +- reuse Qwen3's existing TP runtime +- shard the dense-compatible Qwen3.5 pieces first +- defer true linear-attention/GDR TP until the dense TP path is correct +- gate the work with Qwen3.5 HF logits and scheduler tests + +## Alignment With Qwen3 TP + +Qwen3 already has the main tensor-parallel runtime skeleton. Qwen3.5 should not invent a second parallel runtime. + +Reuse from Qwen3: + +- controller/worker broadcast execution model +- `RequestId` as request identity +- coarse-grained step protocol +- rank-local worker-owned model state +- rank-local CUDA context, cuBLAS, and graph resources +- NCCL hidden all-reduce wrapper +- tensor-parallel config validation pattern +- full-attention head sharding pattern +- MLP intermediate sharding pattern +- first-pass replicated embedding/lm_head simplification + +The Qwen3.5 design should focus on what differs from Qwen3: + +- which Qwen3.5 components can directly reuse Qwen3 TP +- which components cannot because of hybrid/recurrent state +- why the work is split into phases +- what each phase must prove before the next phase starts + +## Non-Goals + +Global non-goals: + +- do not redesign Qwen3's TP runtime +- do not merge Qwen3 and Qwen3.5 into a generic scheduler +- do not support multi-node TP +- do not support `TP > 2` +- do not add pipeline parallelism +- do not add data parallelism +- do not add vocab-parallel embedding/lm_head +- do not promise first-version performance parity with vLLM +- do not solve Qwen3.5 prefix-cache/recurrent-state snapshotting here + +Phase 1 non-goals: + +- do not shard linear attention +- do not shard GDR recurrent state +- do not change the GDR Triton AOT kernel shape +- do not change linear-attention conv state layout +- do not optimize away replicated linear-attention compute +- do not introduce a new GDR backend + +Phase 2 non-goals: + +- do not change GDR math +- do not all-reduce recurrent state +- do not move linear recurrent state ownership back into the scheduler +- do not require TP1 and TP2 bitwise identity; use the HF logits tolerance/regret gate + +## Why Two Phases + +Qwen3.5 TP has two different complexity classes. + +The first class is dense TP that Qwen3 already solved: + +- full-attention q/k/v/o sharding +- full-attention local-KV heads +- MLP gate/up/down sharding +- hidden all-reduce +- worker-thread CUDA/NCCL runtime +- request step broadcast + +These are engineering integration tasks for Qwen3.5. + +The second class is Qwen3.5-specific linear attention and GDR state: + +- recurrent state is long-lived request state, not a temporary tensor +- conv state must follow request identity and decode-slot movement +- GDR state must not be all-reduced +- current GDR AOT kernels are built for global value-head shape +- slot compaction, `DropRequest`, and CUDA graph padding must all work with sharded recurrent state + +Doing both classes in one change would make failures hard to localize. A bad output could come from the TP runtime, full-attention sharding, MLP sharding, GDR kernel shape, recurrent state movement, or slot compaction. + +The two-phase plan narrows the debugging surface: + +- Phase 1 proves Qwen3.5 can run correctly under the TP2 runtime with dense TP and replicated linear attention. +- Phase 2 starts from that stable TP2 base and only changes the linear-attention/GDR ownership contract. + +## Qwen3.5 Architecture Summary + +Qwen3.5-4B is a hybrid decoder: + +- 32 layers +- 24 linear-attention layers +- 8 full-attention layers +- full-attention layer indices: `3, 7, 11, 15, 19, 23, 27, 31` +- `hidden_size = 2560` +- `intermediate_size = 9216` +- tied embedding/lm_head +- `vocab_size = 248320` + +Full attention: + +- `num_attention_heads = 16` +- `num_key_value_heads = 4` +- `head_dim = 256` +- `q_dim = 4096` +- `kv_dim = 1024` +- q projection includes an output gate, so q projection output dim is `8192` + +Linear attention: + +- `linear_num_key_heads = 16` +- `linear_key_head_dim = 128` +- `linear_num_value_heads = 32` +- `linear_value_head_dim = 128` +- `qkv_dim = 8192` +- `z_dim = 4096` +- recurrent state per linear layer: `[32, 128, 128] f32` +- conv state per linear layer: `8192 * (conv_kernel_dim - 1)` bf16 + +## Phase 1: Dense TP With Replicated Linear Attention + +### Scope + +Phase 1 reuses the Qwen3 TP runtime and shards only the Qwen3.5 pieces that match standard dense TP. + +Shard: + +- full-attention `q_proj` +- full-attention `k_proj` +- full-attention `v_proj` +- full-attention `o_proj` +- full-attention KV cache +- MLP `gate_proj` +- MLP `up_proj` +- MLP `down_proj` + +Replicate: + +- embedding +- lm_head / tied `embed_tokens` +- all linear-attention weights +- all linear-attention conv state +- all GDR recurrent state +- linear-attention GDR kernels and scratch shape + +### Partition Contract + +Full attention TP2: + +- global q heads: `16` +- local q heads: `8` +- global KV heads: `4` +- local KV heads: `2` +- global q dim: `4096` +- local q dim: `2048` +- global KV dim: `1024` +- local KV dim: `512` +- global gated q projection dim: `8192` +- local gated q projection dim: `4096` + +MLP TP2: + +- global intermediate: `9216` +- local intermediate: `4608` +- local gate/up rows: `4608` each +- local fused gate_up rows: `9216` +- local down input cols: `4608` + +### Execution Semantics + +For full-attention layers: + +1. each rank computes local q/k/v +2. each rank writes local KV heads +3. each rank runs local attention +4. each rank runs local `o_proj` +5. hidden partials are summed with all-reduce + +For MLP: + +1. each rank computes local gate/up +2. each rank computes local activation +3. each rank runs local `down_proj` +4. hidden partials are summed with all-reduce + +For linear-attention layers: + +1. every rank runs the full linear-attention layer +2. every rank updates a full local copy of recurrent state +3. every rank gets the same full hidden output +4. no linear-attention all-reduce is performed + +Important invariant: + +- replicated linear-attention output must not be accidentally summed across ranks + +### Implementation Tasks + +- Add a Qwen3.5 `TensorParallelConfig`. +- Validate TP2 divisibility for full-attention heads, KV heads, and MLP intermediate. +- Let Qwen3.5 engine accept `tp_size=2`. +- Reuse Qwen3's controller/worker execution model. +- Load rank-local Qwen3.5 models. +- Add full-attention row/column shard loading. +- Add MLP row/column shard loading. +- Allocate full-attention KV pools with local KV heads. +- Add all-reduce after full-attention `o_proj`. +- Add all-reduce after MLP `down_proj`. +- Keep linear-attention loader and forward path replicated. +- Preserve `TP=1` behavior. + +### Acceptance Criteria + +Functional: + +- `TP=1` behavior is unchanged +- `TP=2` starts successfully on two local GPUs +- Qwen3.5 HF logits gate passes under TP2 +- Qwen3.5 scheduler e2e passes under TP2 +- a basic HTTP serving smoke test passes under TP2 +- active request finish/drop does not leak worker-local state + +Correctness: + +- TP2 logits match HF golden within existing Qwen3.5 tolerance +- TP2 handles sequential replay +- TP2 handles graph-padded decode bucket replay +- TP2 handles slot compaction replay +- TP2 long-prompt replay passes when the long fixture is available + +Operational: + +- no deadlock on repeated requests +- no CUDA context mismatch +- no cuBLAS handle cross-device issue +- no NCCL hang on normal shutdown +- unsupported TP sizes fail closed + +Out of scope for acceptance: + +- TP2 being faster than TP1 +- TP2 matching vLLM throughput +- memory optimality from sharding linear attention + +## Phase 2: Linear Attention / GDR TP + +### Scope + +Phase 2 turns Qwen3.5 linear attention from replicated execution into true tensor-parallel execution. + +Shard: + +- `in_proj_qkv` +- `in_proj_z` +- `in_proj_b` +- `in_proj_a` +- `dt_bias` +- `A_log` +- conv state +- GDR recurrent state +- linear-attention `out_proj` + +### vLLM Reference + +vLLM's `Qwen3NextForCausalLM` and `QwenGatedDeltaNetAttention` are the primary external reference for Phase 2. + +The relevant contract to mirror, not copy mechanically: + +- GDN state shape depends on `tp_size`. +- q/k/v/z projections are tensor-parallel column projections. +- `out_proj` is row-parallel and reduces back to full hidden. +- `dt_bias` and `A_log` are sharded over local value heads. +- b/a projections are local-value-head aware; some quantized paths may replicate a small projection and slice locally. +- GDR prefill/decode kernels consume local head/state shapes. + +OpenInfer should translate this into its Rust/CUDA execution model: + +- worker-owned rank-local recurrent state +- explicit `RequestId` lifecycle +- local GDR state movement during slot compaction +- fail-closed kernel shape validation + +The reference does not remove the need for OpenInfer-specific correctness gates. It only proves the partition contract is a known working direction. + +### Partition Contract + +TP2 local linear-attention dims: + +- local key heads: `8` +- local value heads: `16` +- local q dim: `1024` +- local k dim: `1024` +- local v dim: `2048` +- local qkv dim: `4096` +- local z dim: `2048` +- local recurrent state: `[16, 128, 128] f32` +- local conv state: `4096 * (conv_kernel_dim - 1)` bf16 + +### Execution Semantics + +For linear-attention layers: + +1. each rank computes local q/k/v/z/b/a +2. each rank updates only local conv state +3. each rank updates only local GDR recurrent state +4. each rank runs local gated RMSNorm/output-gate path +5. each rank runs local `out_proj` +6. hidden partials are summed with all-reduce + +Important invariants: + +- GDR recurrent state is sharded, not replicated +- GDR recurrent state is never all-reduced +- conv state is sharded by local qkv channels +- all-reduce happens after `out_proj`, not before +- request identity is `RequestId`, not a long-lived batch slot + +### Implementation Tasks + +- Add local linear dim helpers to Qwen3.5 config. +- Add sharded loading for `in_proj_qkv`. +- Add sharded loading for `in_proj_z`. +- Add sharded loading for `in_proj_b` / `in_proj_a`. +- Add sharded loading for `dt_bias` / `A_log`. +- Decide and document `norm_weight` behavior. +- Allocate per-rank local recurrent state. +- Allocate per-rank local conv state. +- Update prefill GDR scratch sizes to local heads. +- Add or generalize GDR prefill kernels for local `num_value_heads`. +- Add or generalize GDR decode kernels for local `num_value_heads`. +- Update slot compaction to move local recurrent/conv state. +- Update `DropRequest` cleanup on all ranks. +- Insert all-reduce after linear-attention `out_proj`. + +### Acceptance Criteria + +Functional: + +- Phase 1 TP2 gates still pass +- linear attention is no longer replicated in TP2 +- per-rank recurrent state size is local, not global +- per-rank conv state size is local, not global +- GDR kernels run with local value heads + +Correctness: + +- HF logits gate passes under TP2 +- long HF logits replay passes +- slot compaction replay passes +- prefill chunking and decode recurrence remain consistent +- TP2 greedy output passes existing regret/logprob tolerance +- TP1 remains unchanged + +Operational: + +- no recurrent state leak after request finish/drop +- no stale state after slot reuse +- no CUDA graph padding corruption +- no GDR scratch OOM regression +- no kernel shape mismatch hidden behind unsafe casts + +Performance sanity: + +- TP2 memory footprint should drop relative to Phase 1 because linear state and some weights are sharded +- TP2 should not catastrophically regress versus Phase 1 on basic decode +- exact throughput target is deferred until correctness is stable + +## Validation Commands + +Phase 1 minimum: + +```bash +OPENINFER_CUDA_SM=120 \ +OPENINFER_TEST_MODEL_PATH=/abs/path/to/Qwen3.5-4B \ +cargo test --release --features qwen35-4b \ + -p openinfer-qwen35-4b --test hf_golden_gate -- --nocapture +``` + +```bash +OPENINFER_CUDA_SM=120 \ +OPENINFER_TEST_MODEL_PATH=/abs/path/to/Qwen3.5-4B \ +cargo test --release --features qwen35-4b \ + -p openinfer-qwen35-4b --test e2e_scheduler -- --nocapture +``` + +Add TP2 variants once the CLI/API exists. + +## References + +- `docs/models/qwen3/tp-design.md` +- `openinfer-qwen3-4b/src/config.rs` +- `openinfer-qwen3-4b/src/executor.rs` +- `openinfer-qwen35-4b/src/config.rs` +- `openinfer-qwen35-4b/src/weights.rs` +- `openinfer-qwen35-4b/src/recurrent_state.rs` +- `openinfer-qwen35-4b/src/batch_decode.rs` +- vLLM `Qwen3NextForCausalLM` +- vLLM `QwenGatedDeltaNetAttention` + From aff47fd9ebdd5e0b98dbaf7a1d061e9ff32d2ad7 Mon Sep 17 00:00:00 2001 From: Troll863 Date: Wed, 24 Jun 2026 21:38:08 +0800 Subject: [PATCH 2/4] doc(qwen35): removed placeholder commands and clarify TP q-gate sharding --- docs/models/qwen35/tp-design.md | 33 ++++++++++++++++++--------------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index 60c4b92c..9259fd28 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -177,6 +177,12 @@ Full attention TP2: - global gated q projection dim: `8192` - local gated q projection dim: `4096` +Qwen3.5 full-attention `q_proj` needs a model-specific sharding rule. The +kernel consumes `q_full` as head-local q/gate pairs, so TP sharding must keep +each query head's q rows together with that head's gate rows. Do not reuse a +naive contiguous row range from Qwen3 if it can split q rows and their gate rows +across different ranks. + MLP TP2: - global intermediate: `9216` @@ -384,25 +390,23 @@ Performance sanity: - TP2 should not catastrophically regress versus Phase 1 on basic decode - exact throughput target is deferred until correctness is stable -## Validation Commands +## Validation Plan Phase 1 minimum: -```bash -OPENINFER_CUDA_SM=120 \ -OPENINFER_TEST_MODEL_PATH=/abs/path/to/Qwen3.5-4B \ -cargo test --release --features qwen35-4b \ - -p openinfer-qwen35-4b --test hf_golden_gate -- --nocapture -``` +- Qwen3.5 HF logits gate under TP2 +- Qwen3.5 scheduler e2e under TP2 +- basic HTTP serving smoke under TP2 + +Phase 2 minimum: -```bash -OPENINFER_CUDA_SM=120 \ -OPENINFER_TEST_MODEL_PATH=/abs/path/to/Qwen3.5-4B \ -cargo test --release --features qwen35-4b \ - -p openinfer-qwen35-4b --test e2e_scheduler -- --nocapture -``` +- Phase 1 gates still pass +- long HF logits replay under TP2 +- slot compaction replay under TP2 +- recurrent-state cleanup on finish/drop -Add TP2 variants once the CLI/API exists. +Add concrete, verified commands once the TP2 CLI/API exists and the commands have +been run on a GPU host. ## References @@ -415,4 +419,3 @@ Add TP2 variants once the CLI/API exists. - `openinfer-qwen35-4b/src/batch_decode.rs` - vLLM `Qwen3NextForCausalLM` - vLLM `QwenGatedDeltaNetAttention` - From e91b6f936b6b8fee2dc587fcde4c7352b775bdf6 Mon Sep 17 00:00:00 2001 From: Troll863 Date: Thu, 25 Jun 2026 16:48:24 +0800 Subject: [PATCH 3/4] docs(qwen35): tighten TP design around degree-parametric contract --- docs/index.md | 2 +- docs/models/qwen35/tp-design.md | 434 ++++++++------------------------ 2 files changed, 105 insertions(+), 331 deletions(-) diff --git a/docs/index.md b/docs/index.md index 9c7168ce..393f2f53 100644 --- a/docs/index.md +++ b/docs/index.md @@ -47,7 +47,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen35/model-crate.md` | `openinfer-qwen35-4b` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35-4b` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | | `models/qwen35/kernel-plan.md` | Qwen3.5-4B has a `openinfer_qwen35_4b::kernel_plan()` static descriptor mirroring the qwen3 module — enumerates every prefill/decode/unified op with its Rust call site, backend, and notes, so you can dump the active kernel mix without reading call sites. Pure refactor (issue #256), no kernel behavior change. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | -| `models/qwen35/tp-design.md` | Qwen3.5 TP2 design: reuse Qwen3's controller/worker TP runtime; Phase 1 shards full-attn+MLP while replicating linear/GDR, Phase 2 shards linear attention and GDR state using vLLM's Qwen3Next/GDN contract as reference. | +| `models/qwen35/tp-design.md` | Qwen3.5 TP design: reuse Qwen3's controller/worker runtime, keep the partition contract degree-parametric with TP2 as the first validation target, and split dense TP from later linear/GDR sharding. | ## models / deepseek-v4 diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index 9259fd28..bbc1701b 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,116 +1,49 @@ # Qwen3.5 Tensor Parallelism Design -> **TL;DR:** Qwen3.5 TP2 should reuse Qwen3's controller/worker TP runtime. Phase 1 shards full-attention and MLP while keeping linear attention/GDR replicated, so the dense TP path can be validated first; Phase 2 then shards linear attention, conv state, GDR recurrent state, and the GDR kernels using vLLM's Qwen3Next/GDN TP contract as the main reference. +> **TL;DR:** Qwen3.5 tensor parallelism should reuse Qwen3's controller/worker TP runtime and stay degree-parametric. Validate `TP=2` first, fail closed on indivisible degrees, shard dense full-attention/MLP before tackling linear-attention GDR state. > > **Last touched:** 2026-06 ## Goal -Add single-node `TP=2` support for `Qwen3.5-4B`. +Add tensor-parallel support for `Qwen3.5-4B` by reusing the Qwen3 TP runtime instead of designing a second parallel execution stack. -The first goal is correctness and architectural integration, not peak performance. +The implementation should be degree-parametric where the model dimensions divide cleanly. `TP=2` is the first validation target, not an architectural limit. Unsupported or indivisible degrees must fail closed before model load. -Concrete goals: +## Qwen3 Runtime Reuse -- keep `TP=1` healthy -- add a fail-closed `TP=2` path -- reuse Qwen3's existing TP runtime -- shard the dense-compatible Qwen3.5 pieces first -- defer true linear-attention/GDR TP until the dense TP path is correct -- gate the work with Qwen3.5 HF logits and scheduler tests - -## Alignment With Qwen3 TP - -Qwen3 already has the main tensor-parallel runtime skeleton. Qwen3.5 should not invent a second parallel runtime. - -Reuse from Qwen3: +Reuse the Qwen3 TP shape: - controller/worker broadcast execution model -- `RequestId` as request identity -- coarse-grained step protocol +- `RequestId` request identity +- coarse-grained prefill/decode/unified/drop step protocol - rank-local worker-owned model state -- rank-local CUDA context, cuBLAS, and graph resources -- NCCL hidden all-reduce wrapper -- tensor-parallel config validation pattern -- full-attention head sharding pattern -- MLP intermediate sharding pattern -- first-pass replicated embedding/lm_head simplification - -The Qwen3.5 design should focus on what differs from Qwen3: - -- which Qwen3.5 components can directly reuse Qwen3 TP -- which components cannot because of hybrid/recurrent state -- why the work is split into phases -- what each phase must prove before the next phase starts - -## Non-Goals - -Global non-goals: - -- do not redesign Qwen3's TP runtime -- do not merge Qwen3 and Qwen3.5 into a generic scheduler -- do not support multi-node TP -- do not support `TP > 2` -- do not add pipeline parallelism -- do not add data parallelism -- do not add vocab-parallel embedding/lm_head -- do not promise first-version performance parity with vLLM -- do not solve Qwen3.5 prefix-cache/recurrent-state snapshotting here +- rank-local CUDA context, cuBLAS, graph, and NCCL resources +- hidden all-reduce after row-parallel projections +- replicated embedding/lm_head as the first-pass simplification -Phase 1 non-goals: +Qwen3.5-specific design work should stay focused on model geometry and state ownership: hybrid layer layout, gated q projection, linear-attention conv state, and GDR recurrent state. -- do not shard linear attention -- do not shard GDR recurrent state -- do not change the GDR Triton AOT kernel shape -- do not change linear-attention conv state layout -- do not optimize away replicated linear-attention compute -- do not introduce a new GDR backend +## Boundaries -Phase 2 non-goals: +This design does not cover multi-node TP, data parallelism, pipeline parallelism, vocab-parallel embedding/lm_head, or Qwen3.5 prefix-cache/recurrent-state snapshots. -- do not change GDR math -- do not all-reduce recurrent state -- do not move linear recurrent state ownership back into the scheduler -- do not require TP1 and TP2 bitwise identity; use the HF logits tolerance/regret gate +Phase 1 does not shard linear attention or change GDR kernel shapes. Phase 2 does not change GDR math, does not all-reduce recurrent state, and does not move recurrent state ownership back into the scheduler. -## Why Two Phases +## Why Dense First, GDR Second -Qwen3.5 TP has two different complexity classes. +Qwen3.5 has two separable TP problems. -The first class is dense TP that Qwen3 already solved: +The dense part is already proven by Qwen3: full-attention head sharding, local KV heads, MLP intermediate sharding, all-reduce after row-parallel projections, and worker-thread CUDA/NCCL execution. -- full-attention q/k/v/o sharding -- full-attention local-KV heads -- MLP gate/up/down sharding -- hidden all-reduce -- worker-thread CUDA/NCCL runtime -- request step broadcast +The linear-attention part is Qwen3.5-specific: conv state and GDR recurrent state are long-lived request state, current GDR AOT kernels are built for the global value-head shape, and slot compaction / graph padding / `DropRequest` must all preserve rank-local recurrent state. If dense TP and GDR TP land together, failures are hard to attribute. Phase 1 narrows correctness debugging to runtime + dense sharding; Phase 2 then isolates the GDR/recurrent contract. -These are engineering integration tasks for Qwen3.5. +## Architecture Summary -The second class is Qwen3.5-specific linear attention and GDR state: +Qwen3.5-4B: -- recurrent state is long-lived request state, not a temporary tensor -- conv state must follow request identity and decode-slot movement -- GDR state must not be all-reduced -- current GDR AOT kernels are built for global value-head shape -- slot compaction, `DropRequest`, and CUDA graph padding must all work with sharded recurrent state - -Doing both classes in one change would make failures hard to localize. A bad output could come from the TP runtime, full-attention sharding, MLP sharding, GDR kernel shape, recurrent state movement, or slot compaction. - -The two-phase plan narrows the debugging surface: - -- Phase 1 proves Qwen3.5 can run correctly under the TP2 runtime with dense TP and replicated linear attention. -- Phase 2 starts from that stable TP2 base and only changes the linear-attention/GDR ownership contract. - -## Qwen3.5 Architecture Summary - -Qwen3.5-4B is a hybrid decoder: - -- 32 layers -- 24 linear-attention layers -- 8 full-attention layers -- full-attention layer indices: `3, 7, 11, 15, 19, 23, 27, 31` +- 32 layers: 24 linear attention + 8 full attention +- full-attention layers: `3, 7, 11, 15, 19, 23, 27, 31` - `hidden_size = 2560` - `intermediate_size = 9216` - tied embedding/lm_head @@ -121,9 +54,9 @@ Full attention: - `num_attention_heads = 16` - `num_key_value_heads = 4` - `head_dim = 256` -- `q_dim = 4096` -- `kv_dim = 1024` -- q projection includes an output gate, so q projection output dim is `8192` +- `q_dim = num_attention_heads * head_dim = 4096` +- `kv_dim = num_key_value_heads * head_dim = 1024` +- q projection includes an output gate, so gated q projection output dim is `2 * q_dim = 8192` Linear attention: @@ -131,282 +64,123 @@ Linear attention: - `linear_key_head_dim = 128` - `linear_num_value_heads = 32` - `linear_value_head_dim = 128` -- `qkv_dim = 8192` -- `z_dim = 4096` -- recurrent state per linear layer: `[32, 128, 128] f32` -- conv state per linear layer: `8192 * (conv_kernel_dim - 1)` bf16 - -## Phase 1: Dense TP With Replicated Linear Attention - -### Scope - -Phase 1 reuses the Qwen3 TP runtime and shards only the Qwen3.5 pieces that match standard dense TP. - -Shard: - -- full-attention `q_proj` -- full-attention `k_proj` -- full-attention `v_proj` -- full-attention `o_proj` -- full-attention KV cache -- MLP `gate_proj` -- MLP `up_proj` -- MLP `down_proj` - -Replicate: - -- embedding -- lm_head / tied `embed_tokens` -- all linear-attention weights -- all linear-attention conv state -- all GDR recurrent state -- linear-attention GDR kernels and scratch shape - -### Partition Contract - -Full attention TP2: - -- global q heads: `16` -- local q heads: `8` -- global KV heads: `4` -- local KV heads: `2` -- global q dim: `4096` -- local q dim: `2048` -- global KV dim: `1024` -- local KV dim: `512` -- global gated q projection dim: `8192` -- local gated q projection dim: `4096` - -Qwen3.5 full-attention `q_proj` needs a model-specific sharding rule. The -kernel consumes `q_full` as head-local q/gate pairs, so TP sharding must keep -each query head's q rows together with that head's gate rows. Do not reuse a -naive contiguous row range from Qwen3 if it can split q rows and their gate rows -across different ranks. - -MLP TP2: - -- global intermediate: `9216` -- local intermediate: `4608` -- local gate/up rows: `4608` each -- local fused gate_up rows: `9216` -- local down input cols: `4608` +- `linear_q_dim = linear_num_key_heads * linear_key_head_dim = 2048` +- `linear_k_dim = linear_q_dim` +- `linear_v_dim = linear_num_value_heads * linear_value_head_dim = 4096` +- `linear_qkv_dim = linear_q_dim + linear_k_dim + linear_v_dim = 8192` +- `linear_z_dim = linear_v_dim = 4096` +- recurrent state per linear layer: `[linear_num_value_heads, linear_key_head_dim, linear_value_head_dim] f32` +- conv state per linear layer: `linear_qkv_dim * (conv_kernel_dim - 1)` bf16 -### Execution Semantics +## Partition Contract -For full-attention layers: +For any candidate `tp`, require: -1. each rank computes local q/k/v -2. each rank writes local KV heads -3. each rank runs local attention -4. each rank runs local `o_proj` -5. hidden partials are summed with all-reduce +- `num_attention_heads % tp == 0` +- `num_key_value_heads % tp == 0` +- `intermediate_size % tp == 0` +- Phase 2 additionally requires `linear_num_key_heads % tp == 0` and `linear_num_value_heads % tp == 0` -For MLP: +Full attention local dimensions: -1. each rank computes local gate/up -2. each rank computes local activation -3. each rank runs local `down_proj` -4. hidden partials are summed with all-reduce +- `local_q_heads = num_attention_heads / tp` +- `local_kv_heads = num_key_value_heads / tp` +- `local_q_dim = local_q_heads * head_dim` +- `local_kv_dim = local_kv_heads * head_dim` +- `local_gated_q_dim = 2 * local_q_dim` -For linear-attention layers: +Qwen3.5 full-attention `q_proj` must be sharded by head-local q/gate pairs. Each rank owns a contiguous query-head range, and for each owned head it must receive both that head's q rows and that head's gate rows. Do not reuse a naive contiguous row shard if the physical layout can split q rows from their gate rows. -1. every rank runs the full linear-attention layer -2. every rank updates a full local copy of recurrent state -3. every rank gets the same full hidden output -4. no linear-attention all-reduce is performed +MLP local dimensions: -Important invariant: +- `local_intermediate = intermediate_size / tp` +- local fused `gate_up_proj` rows: `2 * local_intermediate` +- local `down_proj` input cols: `local_intermediate` -- replicated linear-attention output must not be accidentally summed across ranks +Linear-attention local dimensions for Phase 2: -### Implementation Tasks +- `local_linear_key_heads = linear_num_key_heads / tp` +- `local_linear_value_heads = linear_num_value_heads / tp` +- `local_linear_q_dim = local_linear_key_heads * linear_key_head_dim` +- `local_linear_k_dim = local_linear_q_dim` +- `local_linear_v_dim = local_linear_value_heads * linear_value_head_dim` +- `local_linear_qkv_dim = local_linear_q_dim + local_linear_k_dim + local_linear_v_dim` +- `local_linear_z_dim = local_linear_v_dim` +- local recurrent state: `[local_linear_value_heads, linear_key_head_dim, linear_value_head_dim] f32` +- local conv state: `local_linear_qkv_dim * (conv_kernel_dim - 1)` bf16 -- Add a Qwen3.5 `TensorParallelConfig`. -- Validate TP2 divisibility for full-attention heads, KV heads, and MLP intermediate. -- Let Qwen3.5 engine accept `tp_size=2`. -- Reuse Qwen3's controller/worker execution model. -- Load rank-local Qwen3.5 models. -- Add full-attention row/column shard loading. -- Add MLP row/column shard loading. -- Allocate full-attention KV pools with local KV heads. -- Add all-reduce after full-attention `o_proj`. -- Add all-reduce after MLP `down_proj`. -- Keep linear-attention loader and forward path replicated. -- Preserve `TP=1` behavior. +## Phase 1: Dense TP, Replicated Linear Attention -### Acceptance Criteria - -Functional: - -- `TP=1` behavior is unchanged -- `TP=2` starts successfully on two local GPUs -- Qwen3.5 HF logits gate passes under TP2 -- Qwen3.5 scheduler e2e passes under TP2 -- a basic HTTP serving smoke test passes under TP2 -- active request finish/drop does not leak worker-local state +Shard: -Correctness: +- full-attention `q_proj`, `k_proj`, `v_proj`, `o_proj` +- full-attention KV cache over local KV heads +- MLP `gate_proj`, `up_proj`, `down_proj` -- TP2 logits match HF golden within existing Qwen3.5 tolerance -- TP2 handles sequential replay -- TP2 handles graph-padded decode bucket replay -- TP2 handles slot compaction replay -- TP2 long-prompt replay passes when the long fixture is available +Replicate: -Operational: +- embedding and tied lm_head +- all linear-attention weights +- all linear-attention conv state +- all GDR recurrent state +- existing GDR kernels and scratch shapes -- no deadlock on repeated requests -- no CUDA context mismatch -- no cuBLAS handle cross-device issue -- no NCCL hang on normal shutdown -- unsupported TP sizes fail closed +Execution: -Out of scope for acceptance: +- full-attention: local q/k/v + local attention + local `o_proj`, then all-reduce hidden +- MLP: local gate/up + local activation + local `down_proj`, then all-reduce hidden +- linear attention: every rank runs the full layer and updates a full local recurrent-state copy; do not all-reduce replicated linear-attention output -- TP2 being faster than TP1 -- TP2 matching vLLM throughput -- memory optimality from sharding linear attention +Validation scope: -## Phase 2: Linear Attention / GDR TP +- first validated degree: `TP=2` +- Qwen3.5 HF logits gate +- Qwen3.5 scheduler e2e +- basic TP2 serving smoke +- startup fails closed for unsupported or indivisible degrees -### Scope +## Phase 2: Sharded Linear Attention / GDR -Phase 2 turns Qwen3.5 linear attention from replicated execution into true tensor-parallel execution. +Phase 2 converts linear attention from replicated execution to true TP execution. Shard: -- `in_proj_qkv` -- `in_proj_z` -- `in_proj_b` -- `in_proj_a` -- `dt_bias` -- `A_log` +- `in_proj_qkv`, `in_proj_z`, `in_proj_b`, `in_proj_a` +- `dt_bias`, `A_log` - conv state - GDR recurrent state - linear-attention `out_proj` -### vLLM Reference - -vLLM's `Qwen3NextForCausalLM` and `QwenGatedDeltaNetAttention` are the primary external reference for Phase 2. - -The relevant contract to mirror, not copy mechanically: - -- GDN state shape depends on `tp_size`. -- q/k/v/z projections are tensor-parallel column projections. -- `out_proj` is row-parallel and reduces back to full hidden. -- `dt_bias` and `A_log` are sharded over local value heads. -- b/a projections are local-value-head aware; some quantized paths may replicate a small projection and slice locally. -- GDR prefill/decode kernels consume local head/state shapes. - -OpenInfer should translate this into its Rust/CUDA execution model: - -- worker-owned rank-local recurrent state -- explicit `RequestId` lifecycle -- local GDR state movement during slot compaction -- fail-closed kernel shape validation - -The reference does not remove the need for OpenInfer-specific correctness gates. It only proves the partition contract is a known working direction. - -### Partition Contract - -TP2 local linear-attention dims: +Execution: -- local key heads: `8` -- local value heads: `16` -- local q dim: `1024` -- local k dim: `1024` -- local v dim: `2048` -- local qkv dim: `4096` -- local z dim: `2048` -- local recurrent state: `[16, 128, 128] f32` -- local conv state: `4096 * (conv_kernel_dim - 1)` bf16 +- each rank computes local q/k/v/z/b/a +- each rank updates only local conv state and local GDR recurrent state +- each rank runs local gated RMSNorm/output-gate work +- each rank runs local `out_proj` +- all-reduce happens after `out_proj` -### Execution Semantics +Never all-reduce GDR recurrent state or conv state. Their ownership is rank-local and request-local. -For linear-attention layers: - -1. each rank computes local q/k/v/z/b/a -2. each rank updates only local conv state -3. each rank updates only local GDR recurrent state -4. each rank runs local gated RMSNorm/output-gate path -5. each rank runs local `out_proj` -6. hidden partials are summed with all-reduce - -Important invariants: - -- GDR recurrent state is sharded, not replicated -- GDR recurrent state is never all-reduced -- conv state is sharded by local qkv channels -- all-reduce happens after `out_proj`, not before -- request identity is `RequestId`, not a long-lived batch slot - -### Implementation Tasks - -- Add local linear dim helpers to Qwen3.5 config. -- Add sharded loading for `in_proj_qkv`. -- Add sharded loading for `in_proj_z`. -- Add sharded loading for `in_proj_b` / `in_proj_a`. -- Add sharded loading for `dt_bias` / `A_log`. -- Decide and document `norm_weight` behavior. -- Allocate per-rank local recurrent state. -- Allocate per-rank local conv state. -- Update prefill GDR scratch sizes to local heads. -- Add or generalize GDR prefill kernels for local `num_value_heads`. -- Add or generalize GDR decode kernels for local `num_value_heads`. -- Update slot compaction to move local recurrent/conv state. -- Update `DropRequest` cleanup on all ranks. -- Insert all-reduce after linear-attention `out_proj`. - -### Acceptance Criteria - -Functional: - -- Phase 1 TP2 gates still pass -- linear attention is no longer replicated in TP2 -- per-rank recurrent state size is local, not global -- per-rank conv state size is local, not global -- GDR kernels run with local value heads - -Correctness: - -- HF logits gate passes under TP2 -- long HF logits replay passes -- slot compaction replay passes -- prefill chunking and decode recurrence remain consistent -- TP2 greedy output passes existing regret/logprob tolerance -- TP1 remains unchanged - -Operational: - -- no recurrent state leak after request finish/drop -- no stale state after slot reuse -- no CUDA graph padding corruption -- no GDR scratch OOM regression -- no kernel shape mismatch hidden behind unsafe casts - -Performance sanity: - -- TP2 memory footprint should drop relative to Phase 1 because linear state and some weights are sharded -- TP2 should not catastrophically regress versus Phase 1 on basic decode -- exact throughput target is deferred until correctness is stable +### vLLM Reference -## Validation Plan +Use vLLM's `Qwen3NextForCausalLM` / `QwenGatedDeltaNetAttention` as the reference contract, not as code to copy mechanically: -Phase 1 minimum: +- GDN state shape depends on `tp_size` +- q/k/v/z projections are tensor-parallel column projections +- `out_proj` is row-parallel and reduces back to full hidden +- `dt_bias` and `A_log` are sharded over local value heads +- b/a projections are local-value-head aware; some quantized paths may replicate small projections and slice locally +- GDR prefill/decode kernels consume local head/state shapes -- Qwen3.5 HF logits gate under TP2 -- Qwen3.5 scheduler e2e under TP2 -- basic HTTP serving smoke under TP2 +OpenInfer-specific work remains: worker-owned rank-local recurrent state, `RequestId` lifecycle, local-state slot compaction, `DropRequest` cleanup, and fail-closed kernel-shape validation. -Phase 2 minimum: +Validation scope: - Phase 1 gates still pass -- long HF logits replay under TP2 -- slot compaction replay under TP2 +- long HF logits replay under the validated degree +- slot compaction replay - recurrent-state cleanup on finish/drop - -Add concrete, verified commands once the TP2 CLI/API exists and the commands have -been run on a GPU host. +- no stale local recurrent state after slot reuse ## References From a05b2bcae8f449b34f590fedd7856d30ebf39149 Mon Sep 17 00:00:00 2001 From: Mr_troll863 Date: Mon, 29 Jun 2026 15:34:16 +0800 Subject: [PATCH 4/4] docs(qwen35): clarify TP phase 1 contract --- docs/index.md | 4 +-- docs/models/qwen35/roadmap.md | 6 ++-- docs/models/qwen35/tp-design.md | 51 ++++++++++++++++++++++++++++++++- 3 files changed, 55 insertions(+), 6 deletions(-) diff --git a/docs/index.md b/docs/index.md index 393f2f53..fb443ff4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -40,14 +40,14 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | -| `models/qwen35/roadmap.md` | Qwen3.5-4B roadmap (2026-06 review): decode-tuning refresh improves direct TPOT by 2-3%, while vLLM still leads 1024/256 HTTP decode and high-concurrency throughput. Open items: HND prefill staging, prefix-cache design, serving concurrency. | +| `models/qwen35/roadmap.md` | Qwen3.5-4B roadmap (2026-06 review): decode-tuning refresh improves direct TPOT by 2-3%, while vLLM still leads 1024/256 HTTP decode and high-concurrency throughput. Open items: HND prefill staging, TP Phase 1 implementation, prefix-cache design, serving concurrency. | | `models/qwen35/kv-admission.md` | Issue #254 complete: Qwen3.5 now uses full-lifetime KV admission, deferred pressure handling, impossible-request rejection, explicit error semantics, direct rejection-event coverage, RTX 5090 e2e, and real HTTP pressure/post-pressure validation. | | `models/qwen35/optimization.md` | Hybrid 24 linear + 8 full attn optimization ledger. Decode-tuning refresh fuses MLP gate/up and tunes decode cublasLt buckets, improving direct TPOT by 2-3%; vLLM still leads 1024/256 HTTP decode. | | `models/qwen35/accuracy.md` | Qwen3.5-4B HF bf16 logits goldens through `past_key_values`: short replay covers sequential graph, bucket-straddling batched graph, and slot-compaction; long replay covers 4097/8192-token prompts; full GSM8K 8-shot now matches the HF baseline within 0.15 percentage points. | | `models/qwen35/model-crate.md` | `openinfer-qwen35-4b` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35-4b` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | | `models/qwen35/kernel-plan.md` | Qwen3.5-4B has a `openinfer_qwen35_4b::kernel_plan()` static descriptor mirroring the qwen3 module — enumerates every prefill/decode/unified op with its Rust call site, backend, and notes, so you can dump the active kernel mix without reading call sites. Pure refactor (issue #256), no kernel behavior change. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | -| `models/qwen35/tp-design.md` | Qwen3.5 TP design: reuse Qwen3's controller/worker runtime, keep the partition contract degree-parametric with TP2 as the first validation target, and split dense TP from later linear/GDR sharding. | +| `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 is eager dense TP on Qwen3's controller/worker runtime; validate TP2 first, fail closed for indivisible degrees and TP+CUDA Graph, shard dense full-attention/MLP, and leave sharded linear/GDR state to follow-up. | ## models / deepseek-v4 diff --git a/docs/models/qwen35/roadmap.md b/docs/models/qwen35/roadmap.md index b8950bd3..f60bde9d 100644 --- a/docs/models/qwen35/roadmap.md +++ b/docs/models/qwen35/roadmap.md @@ -1,6 +1,6 @@ # Qwen3.5-4B Roadmap -> **TL;DR:** Qwen3.5-4B is decode-correct and still improving: the decode-tuning refresh improves direct TPOT by `2.1-3.2%`, while vLLM still leads 1024/256 HTTP decode and high-concurrency throughput. Long-prompt HF logits and GSM8K gates cover the old 4096-position RoPE boundary. Remaining structural items are HND prefill staging, prefix-cache design, and the serving-level concurrency gap. +> **TL;DR:** Qwen3.5-4B is decode-correct and still improving: the decode-tuning refresh improves direct TPOT by `2.1-3.2%`, while vLLM still leads 1024/256 HTTP decode and high-concurrency throughput. Long-prompt HF logits and GSM8K gates cover the old 4096-position RoPE boundary. Remaining structural items are HND prefill staging, TP Phase 1 implementation, prefix-cache design, and the serving-level concurrency gap. > > **Last touched:** 2026-06 @@ -20,7 +20,7 @@ Tracking issue: see the `[Model] Qwen3.5-4B roadmap` GitHub issue. Sibling doc: | Admission | ✓ existing full-lifetime KV admission and explicit `Rejected` events cover impossible KV requests; #253 adds the context-window rejection reason before prefill/decode | `scheduler.rs`, `src/scheduler/plan.rs`, `docs/models/qwen35/kv-admission.md` | | Scheduler tests | Partial: current plan selection, full-lifetime admission, context-window rejection, slot assignment, and slot-compaction decisions are CPU-tested; GPU execution remains coupled to the production scheduler | `src/scheduler/plan.rs` | | Step tail | Local branch verified: #353 batches the prefill final norm/lm_head tail, samples decode/unified rows from batched logits, and keeps host full-vocab copies only for requested logprobs; HF/e2e gates pass, short-output serving A/B shows TTFT benefit, long-decode TPOT remains a no-claim diagnostic | `docs/models/qwen35/batched-step-tail.md` | -| TP | ✗ absent (single GPU only) | — | +| TP | Design settled for Phase 1 eager dense TP: reuse Qwen3 controller/worker runtime, validate `TP=2` first, shard dense full-attention/MLP, replicate linear-attention/GDR state per rank, and fail closed for `TP > 1` CUDA Graph or indivisible shapes. Implementation still pending. | `docs/models/qwen35/tp-design.md` | | Prefix cache | ✗ absent; recurrent GDR state (~48MB per boundary snapshot) makes "prefix hit" itself a design question | — | ## Roadmap @@ -41,7 +41,7 @@ Tracking issue: see the `[Model] Qwen3.5-4B roadmap` GitHub issue. Sibling doc: ### Later -- **TP** — no sharding design exists for the hybrid stack (GDR state sharding is the open question). Design-first, no driver today. +- **TP Phase 2** — sharded linear attention / GDR recurrent state remains future design work after Phase 1 eager dense TP is correct. - **CUDA-graph prefill** — prefill is eager and serial; revisit after 6 changes the memory layout. ## Cleanup ledger diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index bbc1701b..210fcd2c 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,6 +1,6 @@ # Qwen3.5 Tensor Parallelism Design -> **TL;DR:** Qwen3.5 tensor parallelism should reuse Qwen3's controller/worker TP runtime and stay degree-parametric. Validate `TP=2` first, fail closed on indivisible degrees, shard dense full-attention/MLP before tackling linear-attention GDR state. +> **TL;DR:** Qwen3.5 tensor parallelism should reuse Qwen3's controller/worker TP runtime and stay degree-parametric. Phase 1 is correctness-first eager dense TP: validate `TP=2` first, fail closed on indivisible degrees and `TP > 1` CUDA Graph, shard dense full-attention/MLP, and keep linear-attention/GDR state replicated per rank before tackling sharded GDR state. > > **Last touched:** 2026-06 @@ -30,6 +30,35 @@ This design does not cover multi-node TP, data parallelism, pipeline parallelism Phase 1 does not shard linear attention or change GDR kernel shapes. Phase 2 does not change GDR math, does not all-reduce recurrent state, and does not move recurrent state ownership back into the scheduler. +## Settled Phase 1 Contract + +These decisions are settled before implementation starts. + +- `TP=1` must preserve the current single-GPU behavior. +- `TP=2` is the first correctness target. The implementation may stay degree-parametric, but unsupported or indivisible degrees must fail before model load. +- `TP > 1` is eager-only in Phase 1. CUDA Graph under TP must fail closed instead of silently falling back or partially capturing. +- Reuse the Qwen3 controller/worker broadcast execution model and avoid a second long-lived Qwen3.5-specific TP runtime shape. +- Shard dense full-attention and MLP operators. +- Replicate embedding and tied `lm_head`. +- Replicate linear-attention/GDR weights in Phase 1. +- Each rank worker owns and mutates its own full linear-attention conv state and GDR recurrent state copy. +- The scheduler owns logical request lifecycle and logical KV/page lifecycle only. +- Full-attention KV is physically rank-local and sharded by local KV heads, but one logical request/page assignment is mirrored across all ranks. +- `DropRequest`, finish cleanup, cancellation cleanup, and slot reuse must release or reset the corresponding rank-local KV/recurrent/conv state on every rank. +- Qwen3.5 gated `q_proj` slicing is an explicit acceptance gate: every rank must receive both q rows and gate rows for its local query heads. +- MLP gate/up row sharding and down column sharding require explicit reconstruction or layout tests. + +## Still Open / Future Discussion + +These topics should not block Phase 1 eager dense TP, but they remain design work before any later implementation. + +- TP CUDA Graph support: graph state ownership per rank, synchronized capture/replay order, NCCL capture behavior, graph padding slots, and recurrent/conv D2D slot compaction under capture. +- Sharded linear-attention/GDR execution: local GDR AOT kernel shapes, local recurrent-state layout, local conv state layout, and Phase 2 weight slicing. +- TP-aware prefix cache or recurrent-state snapshots. +- Vocab-parallel embedding or `lm_head`. +- Multi-node TP, data parallelism, and pipeline parallelism. +- Performance optimization claims. Phase 1 is a correctness/runtime milestone, not a throughput milestone. + ## Why Dense First, GDR Second Qwen3.5 has two separable TP problems. @@ -131,13 +160,33 @@ Execution: - MLP: local gate/up + local activation + local `down_proj`, then all-reduce hidden - linear attention: every rank runs the full layer and updates a full local recurrent-state copy; do not all-reduce replicated linear-attention output +State ownership: + +- scheduler owns request admission, request identity, logical page allocation, streaming handles, sampling params, generation counters, and finish bookkeeping +- rank workers own rank-local model shards, rank-local physical KV buffers, rank-local decode buffers, and rank-local recurrent/conv state +- rank 0 is not special for state mutation; it follows the same worker command protocol as other ranks +- non-primary workers may return acknowledgement or step failure only, while the primary worker returns artifacts for scheduler-side result resolution +- all workers must observe the same ordered `RunPrefillStep`, `RunDecodeStep`, `RunUnifiedStep`, `DropRequest`, and `Shutdown` commands + +CUDA Graph: + +- Phase 1 TP execution is eager-only +- `tp_size > 1` with CUDA Graph enabled must return an explicit startup/configuration error before serving requests +- TP graph capture is a follow-up because Qwen3.5 graph state includes recurrent slots, slot compaction, padding slots, and NCCL ordering questions + Validation scope: - first validated degree: `TP=2` - Qwen3.5 HF logits gate - Qwen3.5 scheduler e2e +- long prompt / chunked prefill path +- slot-compaction replay +- finish/drop followed by slot reuse without stale recurrent or conv state +- gated `q_proj` head-local q/gate slicing test +- MLP gate/up shard and down shard reconstruction/layout test - basic TP2 serving smoke - startup fails closed for unsupported or indivisible degrees +- startup fails closed for `tp_size > 1` with CUDA Graph enabled ## Phase 2: Sharded Linear Attention / GDR