RFC(qwen35): TP Phase2

[Mrtroll486]

RFC: Qwen3.5-4B Tensor Parallelism Phase 2

TL;DR: For Qwen3.5-4B, Phase 2 should first add mixed prefill/decode execution to the existing eager TP runtime, then convert the replicated linear-attention/GDR path into true tensor-parallel execution with rank-local weights and recurrent state. Per-device Triton AOT handle management, TP CUDA Graphs, TP-aware prefix caching, vocabulary parallelism, and performance acceptance criteria require separate decisions before they become committed deliverables.

Refs issue #446, #459
Refs pr #627, #450

Motivation

Qwen3.5-4B TP Phase 1 established a correctness-first TP2 serving path by reusing the Qwen3 controller/worker runtime. It shards the dense full-attention and MLP paths, but every rank still loads and executes the complete linear-attention/GDR path and owns a full copy of its conv and recurrent state.

Phase 2 should complete the model-specific part of Qwen3.5-4B tensor parallelism without mixing unrelated features into the same correctness milestone. The work should preserve Phase 1's worker-owned request state and deterministic collective ordering while reducing the replicated linear-attention compute and state on each rank.

Phase 1 Baseline

Phase 1 is complete as a Qwen3.5-4B eager TP2 correctness and runtime milestone.

Implemented:

• TP rank/world-size configuration and dense-dimension divisibility validation.
• Fail-closed startup for TP > 1 with CUDA Graph enabled.
• Rank-local workers owning their CUDA context, cuBLAS handles, NCCL communicator, model shard, physical KV state, conv state, and GDR recurrent state.
• A fail-stop TP execution contract: the first worker failure poisons the executor, response waits and shutdown are bounded, and later operations fail promptly.
• Dense TP weight loading for full-attention q_proj, k_proj, v_proj, o_proj, local KV heads, and MLP gate_proj, up_proj, and down_proj.
• Eager prefill, chunked prefill, eager decode, scheduler integration, and the public multi-device serving path.
• Correct head-interleaved slicing for the full-attention gated q_proj.
• Per-device Triton AOT module/function handles and worker-local NCCL initialization.
• TP2 short and long HF logits gates, scheduler E2E coverage, and OpenAI-compatible HTTP serving smoke tests.
• TP1 short/long HF and scheduler regression coverage after the TP2 changes.

Phase 1 deliberately does not provide:

• TP mixed prefill/decode execution through RunUnifiedStep.
• Sharded linear-attention/GDR weights, kernels, conv state, or recurrent state.
• TP CUDA Graph capture/replay.
• TP-aware prefix caching or recurrent-state snapshots.
• Vocabulary-parallel embedding or lm_head.
• A TP performance claim or numeric performance acceptance gate.

Proposed Phase 2 Scope

Phase 2 is split into two milestones so execution-protocol failures remain distinguishable from model-state and kernel-shape failures.

P2a: TP Mixed Unified Execution

Implement RunUnifiedStep while retaining the Phase 1 replicated linear-attention/GDR path.

Requirements:

• Execute mixed prefill and decode rows in one TP scheduler step.
• Broadcast the same ordered step to every rank.
• Preserve deterministic collective ordering across all ranks, including padded or inactive rows.
• Keep logical request and page accounting in the scheduler while workers continue to own physical KV, conv, and recurrent state.
• Return prefill and decode artifacts from the primary rank without making rank 0 special for model-state mutation.
• Handle normal finish, explicit drop, client disconnect, cancellation, compaction, and slot reuse consistently on every rank.
• Continue to fail closed when TP CUDA Graph is requested.

P2a is a protocol and lifecycle milestone. It must not depend on changing GDR weight or state shapes.

P2b: Sharded Linear Attention and GDR

Convert Qwen3.5-4B's 24 linear-attention layers from replicated execution to true tensor-parallel execution.

Partition Contract

For a candidate TP degree, additionally require:

• linear_num_key_heads % tp_size == 0
• linear_num_value_heads % tp_size == 0

Define rank-local dimensions as:

• local_linear_key_heads = linear_num_key_heads / tp_size
• local_linear_value_heads = linear_num_value_heads / tp_size
• 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

Unsupported or indivisible configurations must fail before model loading.

Weight and State Sharding

Partition the linear-attention projection, parameter, and state surface by the
resolved rank-local key/value-head layout:

• rank-local projection and per-head parameter tensors;
• linear-attention out_proj;
• conv state;
• GDR recurrent state.

The implementation design must map the actual checkpoint packing and kernel
input layout to these local dimensions, then prove the mapping with loader
reconstruction tests. This RFC fixes the rank-local state and collective
contract, not a premature tensor-slice implementation.

Embedding and tied lm_head remain replicated unless vocabulary parallelism is accepted separately.

Execution Contract

Each rank must:

• Compute only its rank-local linear-attention projections and parameters.
• Run GDR prefill and decode kernels against local head and state shapes.
• Update only its rank-local, request-local conv and GDR recurrent state.
• Run local gated RMSNorm and output-gate work.
• Run its row-parallel local out_proj.
• All-reduce the hidden output only after out_proj.

The implementation must never all-reduce conv state or GDR recurrent state. Those states remain rank-local and request-local for their full lifetime.

Kernel Work

• Adapt or regenerate the Triton AOT GDR kernels for local head and recurrent-state shapes.
• Resize rank-local scratch and intermediate buffers to local dimensions.
• Preserve per-device module/function handle ownership for generated AOT stubs.
• Validate kernel-shape support before model loading and fail closed for unsupported combinations.
• Keep collective calls outside GDR state mutation; the only required linear-attention collective is the hidden all-reduce after out_proj.

Capacity and Admission

P2b changes the per-rank recurrent, conv, scratch, and intermediate allocation
shapes. The existing TP fail-stop and rank-safe capacity contract remains a
Phase 2 prerequisite, not a new performance target:

• Recompute the rank-local active-request capacity from the final local state
  and buffer shapes after model/KV loading, with the existing runtime safety
  reserve.
• Advertise and allocate against the minimum safe capacity across ranks; rank
  0 must not advertise capacity that another worker cannot serve.
• Keep full-lifetime KV admission and state-capacity admission consistent so
  expected state or buffer allocation for an admitted request cannot cause one
  rank to OOM while peers enter a collective; unexpected CUDA failures follow
  the fail-stop contract.
• Preserve fail-stop behavior for rank-local execution failure; later executor
  operations must fail promptly rather than leave peers or the controller
  waiting in NCCL.

Request-State Lifecycle

Every worker must apply the same ordered lifecycle transition for a RequestId:

• admission and initial state allocation
• repeated or chunked prefill
• promotion to decode
• mixed-step slot movement
• finish
• explicit drop
• cancellation or client disconnect
• compaction
• slot reuse

Cleanup must release or reset the corresponding local KV, conv, and recurrent state on every rank. Reusing a slot must never expose state from a previous request.

Acceptance Criteria

P2a Gates

• Existing TP1 short/long HF logits gates still pass.
• Existing TP2 short/long HF logits gates still pass.
• Existing TP1 and TP2 scheduler E2E tests still pass.
• Existing TP2 HTTP serving smoke still passes.
• Mixed prefill+decode execution passes under TP2.
• A TP2 mixed-step integration case advances a long request through repeated or
  chunked prefill (base_pos > 0) while another request decodes, then finishes
  or cancels the decode request, compacts slots, and admits a new request into
  the reused slot. It verifies request-to-artifact row mapping and that no
  stale KV, conv, or recurrent state is observable.
• Mixed-step artifacts preserve request identity and the existing greedy,
  sampling, stop, and logprobs semantics for both prefill and decode rows.
• Mixed-step collective ordering remains deterministic when requests finish, cancel, or disconnect.
• TP > 1 with CUDA Graph enabled continues to fail at startup.

P2b Gates

• Linear-attention/GDR loader layout tests reconstruct the expected global tensors from rank-local shards.
• Local GDR kernel/state shape validation rejects unsupported TP degrees before model loading.
• Rank-local GDR weight and live-state allocation checks confirm that P2b does
  not retain an accidental replicated linear-attention path.
• TP capacity accounting tests cover the final local recurrent, conv, scratch,
  and intermediate shapes; the effective max_batch is the minimum safe
  capacity across ranks.
• TP admission pressure defers or rejects work before a rank-local state or
  buffer OOM, and a post-pressure request still completes cleanly.
• TP2 short HF logits replay passes with sharded GDR execution.
• TP2 long HF logits replay passes for the existing 4097- and 8192-token prompts.
• Slot-compaction replay passes with rank-local recurrent and conv state.
• Finish, explicit drop, cancellation, and client-disconnect cleanup pass.
• Slot reuse after every cleanup path shows no stale recurrent or conv state.
• Phase 1 scheduler and HTTP gates pass without reverting to replicated GDR execution.

TP2 on Qwen3.5-4B remains the first required validation degree. The implementation should remain degree-parametric where all required dimensions divide cleanly, but this RFC does not claim untested degrees or other Qwen3.5 sizes as supported.

Non-Goals

• Changing GDR mathematics or introducing a different recurrent algorithm.
• All-reducing, centralizing, or moving conv/GDR recurrent state back into the scheduler.
• Multi-node TP, data parallelism, or pipeline parallelism.
• Silently falling back to replicated linear attention for a requested unsupported TP configuration.
• Treating the unresolved items below as implicit Phase 2 commitments.

Decisions Required

The following items are intentionally unresolved. Each needs an explicit scope and acceptance decision before implementation is included in this RFC.

Per-Device Triton AOT Handle Management

The current per-device Triton AOT patch stores CUmodule and CUfunction
handles in a fixed table indexed by CUDA logical device ordinal. Decide whether
that fixed limit remains a supported runtime boundary or the handle table must
be created for the visible devices on demand.

Questions to resolve:

• If retaining a fixed limit, which logical device ordinals are supported, and
  where does startup reject an unsupported EngineLoadOptions.device_ordinals
  configuration before workers launch kernels?
• If replacing the fixed table, should handles be sized from the visible CUDA
  device count or allocated lazily by current device, and which object owns
  concurrent initialization, failure caching, and teardown?
• How does the contract interact with CUDA_VISIBLE_DEVICES and non-zero
  logical device ordinals selected for TP workers?
• Which smoke coverage proves that every TP worker activates the module and
  function handle for its own device?

If dynamic or demand-driven handle management is accepted, it is a P2a runtime
prerequisite and must land before P2a unified execution. If it is deferred,
P2a must retain an explicit startup-time device-ordinal validation rather than
failing at the first Triton kernel launch.

TP CUDA Graph

Decide whether TP CUDA Graph support belongs in Phase 2, a later phase, or a separate RFC.

Questions to resolve:

• Which object owns graph state on each rank?
• How are capture and replay ordered across ranks?
• Which NCCL operations and communicator setup are capture-safe on the supported stack?
• How are padded graph slots represented across ranks?
• How do recurrent/conv D2D compaction and slot reuse work under capture?
• Does graph support require P2a and P2b to be complete first?
• Which batch buckets and TP degrees are required acceptance targets?

Until decided, TP > 1 with CUDA Graph enabled must continue to fail closed.

TP-Aware Prefix Cache and Recurrent-State Snapshots

Decide whether prefix caching stores only full-attention KV pages or also snapshots linear-attention conv/GDR recurrent state.

Questions to resolve:

• What is the logical cache identity shared across ranks?
• Which physical KV pages are stored per rank?
• What recurrent and conv state is required to resume from a cached prefix without replay?
• Are recurrent snapshots sharded in the same layout as live P2b state?
• How are snapshot lifetime, eviction, cancellation, and partial hits coordinated?
• Is suffix replay acceptable for an initial TP-aware prefix-cache milestone?

Prefix caching and recurrent snapshots are not required for P2a or P2b correctness unless this decision changes the scope.

Vocabulary Parallelism

Decide whether to shard the embedding and tied lm_head instead of retaining Phase 1 replication.

Questions to resolve:

• Is the primary goal HBM reduction, output-tail latency, or support for larger Qwen3.5 variants?
• Should embedding lookup use all-reduce or another ownership scheme?
• Should greedy and sampled token selection operate directly on vocabulary shards?
• How are logprobs and requested top-logprobs reconstructed?
• What communication is required before or after sampling?
• Must tied embedding/lm_head remain physically tied after sharding?

Vocabulary parallelism is not part of P2a or P2b unless separately accepted.

Performance Requirements

Decide whether Phase 2 is only a correctness and memory-layout milestone or must pass numeric performance gates.

Metrics and comparison points to choose:

• Per-rank model-weight and recurrent-state HBM reduction.
• Linear-attention prefill and decode latency before and after P2b.
• End-to-end TTFT, TPOT, ITL, output tokens/s, and peak HBM.
• Solo and concurrent workloads, including mixed prefill/decode traffic.
• Comparison against Phase 1 replicated-GDR TP2, TP1, and a pinned vLLM baseline.
• Allowed regression thresholds for P2a before P2b optimization.
• Minimum speedup or scaling-efficiency requirements for declaring Phase 2 complete.

No performance improvement should be claimed until the workload, hardware, software versions, and thresholds are agreed and measured.

Proposed Delivery Order

1. Resolve per-device Triton AOT handle management; if dynamic or demand-driven
   management is accepted, land it as a P2a runtime prerequisite.
2. Land P2a mixed unified execution with Phase 1 replicated GDR state.
3. Add P2a lifecycle and mixed-step correctness gates.
4. Add P2b config validation and CPU-only shard layout tests.
5. Add rank-local linear-attention weight and state loading.
6. Adapt GDR prefill/decode kernels and scratch buffers to local shapes.
7. Add the hidden all-reduce after local linear-attention out_proj.
8. Run short HF replay, then long replay, then lifecycle and scheduler gates.
9. Resolve the separate CUDA Graph, prefix-cache, vocabulary-parallel, and performance decisions before extending the committed scope.

References

• docs/models/qwen35/tp-design.md
• docs/models/qwen35/tp-implementation.md
• openinfer-qwen35-4b/src/tp_executor.rs
• vLLM Qwen3NextForCausalLM
• vLLM QwenGatedDeltaNetAttention

[xiaguan]

The P2a/P2b split looks right. I suggest we lock Phase 2 to the following scope:

• P2a: implement eager TP RunUnifiedStep with the same ordered prefill/decode plan on every rank. Keep the Phase 1 replicated GDR path. Key lifecycle and artifact mapping by RequestId; do not introduce CUDA Graph slot/compaction semantics into the eager path.
• P2b: shard the linear-attention/GDR weights, kernels, conv state, and recurrent state by local heads. State remains request-local and rank-local; the only linear-attention collective should be the hidden all-reduce after local out_proj.

For the unresolved items:

• Keep the current fixed 16-device Triton AOT handle table and add startup-time validation for unsupported logical ordinals. Dynamic handle allocation is not a P2a prerequisite.
• Keep TP CUDA Graph, TP-aware prefix caching/recurrent snapshots, and vocabulary parallelism in separate follow-up RFCs. TP eager should continue to fail closed when CUDA Graph is requested.
• Treat P2b as a correctness and per-rank HBM-reduction milestone first. It adds one hidden all-reduce for each of the 24 linear-attention layers, so a latency/throughput win is not guaranteed. Please report Phase 1 TP2 vs P2b TP2 A/B data before making a performance claim.

Please deliver P2a and P2b as separate implementation series so protocol/lifecycle failures remain distinguishable from loader/kernel/state-shape failures.

[Mrtroll486]

Thanks, agreed. I’ll update the docs PR to lock this scope, then deliver P2a and P2b as separate PRs. P2a will keep the replicated-GDR eager path; P2b will be treated as a correctness/HBM milestone, with the requested TP2 A/B included before any performance claim.

[github-actions]

This issue has been inactive for 14 days. It will be closed after another 30 days unless there is new activity.