diff --git a/Cargo.lock b/Cargo.lock index 7e94c6f73..691e35c7e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3317,14 +3317,17 @@ dependencies = [ "openinfer-core", "openinfer-kernels", "openinfer-sample", + "openinfer-vllm-frontend", "openinfer-vllm-support", "rand 0.10.1", + "reqwest", "safetensors", "serde", "serde_json", "sha2 0.11.0", "tempfile", "tokio", + "tokio-util", "vllm-text", ] diff --git a/docs/index.md b/docs/index.md index 24892bfd7..aedd8b821 100644 --- a/docs/index.md +++ b/docs/index.md @@ -52,6 +52,8 @@ 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 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/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 implementation record: eager dense TP2 worker/scheduler path, short/long HF logits gates, scheduler e2e, and real OpenAI-compatible HTTP smoke pass; remaining TP work is kept as follow-up, not a Phase 1 claim. | ## models / glm52 diff --git a/docs/models/qwen35/roadmap.md b/docs/models/qwen35/roadmap.md index 3a2dd7749..4f68b55fe 100644 --- a/docs/models/qwen35/roadmap.md +++ b/docs/models/qwen35/roadmap.md @@ -47,7 +47,7 @@ out: | Fault isolation | Open risk: batch-level execution errors can still fail multiple active requests | #654 | | Prefix reuse | Open: bounded joint KV/recurrent/conv snapshot design and implementation | #257 | | DFlash | In flight and opt-in: correctness-first work must stay default-off until gates pass | #434, PR #626, #654 | -| Tensor parallel | Open: engine currently accepts one CUDA device; design-first | #446 | +| Tensor parallel | Phase 1 complete: eager dense TP2 worker/scheduler execution; Phase 2 still needs mixed-step execution and sharded linear-attention/GDR state. | `docs/models/qwen35/tp-implementation.md`, #446 | ## Active Contract diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md new file mode 100644 index 000000000..210fcd2ca --- /dev/null +++ b/docs/models/qwen35/tp-design.md @@ -0,0 +1,244 @@ +# Qwen3.5 Tensor Parallelism Design + +> **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 + +## Goal + +Add tensor-parallel support for `Qwen3.5-4B` by reusing the Qwen3 TP runtime instead of designing a second parallel execution stack. + +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. + +## Qwen3 Runtime Reuse + +Reuse the Qwen3 TP shape: + +- controller/worker broadcast execution model +- `RequestId` request identity +- coarse-grained prefill/decode/unified/drop step protocol +- rank-local worker-owned model state +- 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 + +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. + +## Boundaries + +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. + +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. + +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. + +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. + +## Architecture Summary + +Qwen3.5-4B: + +- 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 +- `vocab_size = 248320` + +Full attention: + +- `num_attention_heads = 16` +- `num_key_value_heads = 4` +- `head_dim = 256` +- `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: + +- `linear_num_key_heads = 16` +- `linear_key_head_dim = 128` +- `linear_num_value_heads = 32` +- `linear_value_head_dim = 128` +- `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 + +## Partition Contract + +For any candidate `tp`, require: + +- `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` + +Full attention local dimensions: + +- `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` + +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. + +MLP local dimensions: + +- `local_intermediate = intermediate_size / tp` +- local fused `gate_up_proj` rows: `2 * local_intermediate` +- local `down_proj` input cols: `local_intermediate` + +Linear-attention local dimensions for Phase 2: + +- `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 + +## Phase 1: Dense TP, Replicated Linear Attention + +Shard: + +- 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` + +Replicate: + +- embedding and tied lm_head +- all linear-attention weights +- all linear-attention conv state +- all GDR recurrent state +- existing GDR kernels and scratch shapes + +Execution: + +- 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 + +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 + +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` +- conv state +- GDR recurrent state +- linear-attention `out_proj` + +Execution: + +- 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` + +Never all-reduce GDR recurrent state or conv state. Their ownership is rank-local and request-local. + +### vLLM Reference + +Use vLLM's `Qwen3NextForCausalLM` / `QwenGatedDeltaNetAttention` as the reference contract, not as code to 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 small projections and slice locally +- GDR prefill/decode kernels consume local head/state shapes + +OpenInfer-specific work remains: worker-owned rank-local recurrent state, `RequestId` lifecycle, local-state slot compaction, `DropRequest` cleanup, and fail-closed kernel-shape validation. + +Validation scope: + +- Phase 1 gates still pass +- long HF logits replay under the validated degree +- slot compaction replay +- recurrent-state cleanup on finish/drop +- no stale local recurrent state after slot reuse + +## 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` diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md new file mode 100644 index 000000000..1df76e5bc --- /dev/null +++ b/docs/models/qwen35/tp-implementation.md @@ -0,0 +1,181 @@ +# Qwen3.5 TP Implementation Record + +> **TL;DR:** Qwen3.5 TP Phase 1 is implemented as correctness-first eager dense TP: TP2 worker/scheduler execution, short/long HF logits gates, scheduler e2e, and real OpenAI-compatible HTTP serving smoke pass. The branch is rebased onto current `main` with the newer engine, sampling, config, and golden-fixture contracts; remaining TP work is tracked as follow-up, not a Phase 1 claim. +> +> **Last touched:** 2026-07 + +## Scope + +This is the implementation record for Qwen3.5 tensor parallelism. The stable architecture contract lives in `docs/models/qwen35/tp-design.md`; this file records what actually landed, what was verified, and what should carry into later phases. + +Keep Phase 1 and Phase 2 in this file until the Phase 2 implementation becomes large enough to split. The same state ownership risks continue across phases, so keeping the history together is useful: Phase 1 proves dense TP and worker-owned request state, while Phase 2 builds on that boundary for mixed prefill/decode and sharded recurrent state. + +Out of scope for this file: + +- local machine paths, NCCL symlink details, and temporary environment setup +- raw command transcripts unless they are part of retained evidence +- benchmark/performance claims + +## Phase 1 Outcome + +Phase 1 is complete as a correctness/runtime milestone. + +Implemented: + +- TP config validation for rank/world size, dense divisibility, and `TP > 1 && CUDA Graph` fail-closed startup. +- Dense TP weight loading for full-attention projections, full-attention KV heads, and MLP projections. +- Rank-local worker executor with worker-owned model shards, KV state, recurrent/conv state, CUDA context, cuBLAS, and NCCL comms. +- Eager TP prefill, chunked prefill, and eager decode. +- Scheduler TP backend that routes chunked prefill and eager decode through TP workers while keeping logical request/page accounting in the scheduler. +- Public multi-device Qwen3.5 engine path and server launch path for `tp_size > 1` with CUDA Graph disabled. +- Real HTTP TP2 serving smoke through the vLLM/OpenAI-compatible frontend. + +Not implemented in Phase 1: + +- TP CUDA Graph capture/replay. +- TP `RunUnifiedStep` mixed prefill+decode execution. +- Sharded linear-attention/GDR weights, kernels, conv state, or recurrent state. +- Vocab-parallel embedding or `lm_head`. +- Prefix-cache or recurrent-state snapshot support. +- Performance claims. + +## Important Fixes + +### Gated q projection layout + +The major numeric blocker was the full-attention gated `q_proj` TP shard layout. + +The wrong assumption was that `q_proj.weight` rows were physically arranged as: + +```text +[all q rows][all gate rows] +``` + +The actual Qwen3.5 kernel contract is per-head interleaved: + +```text +[head0 q][head0 gate][head1 q][head1 gate]... +``` + +For TP2, the fixed loader preserves contiguous head-interleaved ranges: + +- rank 0 loads rows `0..4096` +- rank 1 loads rows `4096..8192` + +The old loader gathered local q rows and local gate rows separately, then rebuilt a `[q][gate]` fused matrix. That corrupted the first full-attention contribution and failed the TP2 HF gate from prefill position `0`. + +### Per-device Triton AOT handles + +Real TP2 prefill exposed that Qwen3.5 GDR Triton AOT C stubs could not cache `CUmodule` / `CUfunction` in process-global state. With two CUDA devices, the rank that loaded a GDR kernel first could leave the other rank with an invalid function handle. + +The generated stubs now cache module/function handles per CUDA device ordinal. This is an implementation constraint worth remembering for future multi-GPU users of generated Triton C stubs. + +Follow-up review tightened this path: the generated stubs now fail closed before indexing the fixed per-device handle tables if `cuCtxGetDevice` returns an ordinal outside the table size. This preserves the Phase 1 static-table implementation while avoiding out-of-bounds writes on high CUDA ordinals. + +### Worker-local NCCL setup + +NCCL comms are initialized inside rank worker threads after each worker binds its CUDA context and initializes thread-local cuBLAS. Creating comms on the controller thread and moving them into workers led to invalid-handle symptoms and hangs. + +This matches the design contract: TP workers own rank-local CUDA/NCCL execution resources. + +### Current-main API compatibility + +Rebasing Phase 1 onto current `main` required preserving the TP execution boundary while adopting newer shared contracts: + +- Hybrid batch decode now builds `Vec<&mut RecurrentState>` from graph-owned slots before entering the common linear-attention helper. This keeps request state in place while satisfying the helper's mutable-reference slice contract. +- `openinfer_sample::select_batch` now requires request-local sampling steps. Phase 1 TP still samples one row at a time and has no request-local sampling counter, so it passes step `0` and retains its existing per-row `sample_seed` offset. Do not substitute batch row indices for request-local steps: that would make seeded output depend on batch composition. +- Qwen3.5 launch and tests use the current `EngineLoadOptions` surface; the removed `enable_prefill_profile` field is no longer supplied. +- TP scheduler tests explicitly set the newer `GenerateRequest::data_parallel_rank` field to `None` because Phase 1 is TP-only, not DP. +- Synthetic TP config/loader fixtures include `tie_word_embeddings`, matching the current `Config35` contract without changing production config loading. +- TP2 short/long HF gates use `Golden::load_for(model_path, long)` and pass the complete `Golden` to metadata validation, matching the model-selected fixture flow used by TP1. + +These are compatibility changes, not extensions of Phase 1 scope. In particular, full seeded-sampling replay under TP should add a request-local completion counter rather than overloading batch position. + +## Validation Evidence + +Phase 1 acceptance coverage: + +- TP2 short HF logits gate passes: + - sequential eager: `108` positions, mean `0.0258`, p99 `0.0801`, max `0.1298` + - batched eager: `72` positions, mean `0.0257`, p99 `0.0809`, max `0.1298` +- TP2 long HF logits gate passes: + - prompts `4097` and `8192`, sequential eager: `18` positions, mean `0.0232`, p99 `0.0792`, max `0.1035` +- TP2 scheduler e2e passes and covers: + - context-window rejection + - greedy/logprobs paths + - sequential requests + - repeated request reuse + - concurrent mixed greedy/sampling requests + - consumer drop + - post-drop scheduler health +- TP2 HTTP serving smoke passes through `openinfer_vllm_frontend::serve`: + - `/v1/models` + - non-streaming `/v1/completions` + - streaming `/v1/completions` + - concurrent completions + - finite logprobs + - chunked prefill forced with `max_prefill_tokens=1` + - `TP2 + CUDA Graph` fail-closed startup +- TP1 regression gates pass after the TP2 additions: + - TP1 short/long HF logits gates + - TP1 scheduler e2e +- Current-main rebase verification passes: + - formatting check + - Qwen3.5 release compilation for all test targets + - `openinfer-server` release compilation with only the `qwen35-4b` model feature + +Known validation constraints: + +- TP2 tests remain ignored by default because they require two CUDA devices, NCCL, and real Qwen3.5 weights. +- Long TP2 HF replay is GPU-memory-sensitive; choose a sufficiently free device pair. +- Qwen3.5 HF golden integration tests should run serially on memory-constrained hosts to avoid unrelated KV-capacity failures from concurrent model loads. + +Stable test knobs: + +- `OPENINFER_TEST_MODEL_PATH`: real Qwen3.5 weights path for HF, scheduler, and serving tests. +- `OPENINFER_TEST_TP_DEVICES`: comma-separated TP2 CUDA ordinals. Defaults to `0,1`; examples: `1,2`, `2,3`. TP2 tests require exactly two distinct ordinals. +- `OPENINFER_TEST_FRONTEND_MODEL_PATH`: optional tokenizer/config metadata path for HTTP serving tests. Defaults to `OPENINFER_TEST_MODEL_PATH` when unset. + +## Follow-Up Work + +The exact Phase 2 split is not decided yet. The items below are retained as follow-up work that should be scoped in the design branch before implementation. + +### TP mixed-step unified execution + +Implement `RunUnifiedStep` under TP while keeping Phase 1's replicated linear-attention/GDR state unless the design branch decides otherwise. + +Goals: + +- Support mixed prefill+decode scheduler steps under TP. +- Preserve deterministic collective ordering across ranks. +- Return mixed prefill/decode artifacts from the primary rank. +- Validate finish/drop/client-disconnect cleanup under mixed-step execution. +- Keep TP CUDA Graph disabled unless a separate graph design is completed. + +Why this should be separated from GDR sharding: + +- Mixed-step scheduling is an execution-protocol problem. +- Sharded linear-attention/GDR is a model-state-shape problem. +- Combining them would make failures hard to attribute. + +### Sharded linear-attention/GDR state + +Shard the Qwen3.5 linear-attention/GDR path after the mixed-step and state-lifecycle contract is clear. + +Expected work: + +- shard linear-attention projection weights +- shard conv state and GDR recurrent state by local value/key heads +- adapt or regenerate GDR kernels for local state shapes +- keep recurrent/conv state rank-local and request-local +- all-reduce only after local linear-attention `out_proj` + +Non-negotiable invariant: + +- Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state. + +## Follow-Ups + +- Promote any stable contract changes discovered here back into `tp-design.md` through the design-doc branch. +- Decide whether Qwen3.5 server CLI should accept arbitrary TP device ordinals instead of only `0..tp_size`. +- Consider lifting the per-device Triton AOT handle lesson into a kernels or runtime subsystem doc if another model hits the same issue. diff --git a/openinfer-kernels/build.rs b/openinfer-kernels/build.rs index aa0207d70..af96b05b0 100644 --- a/openinfer-kernels/build.rs +++ b/openinfer-kernels/build.rs @@ -1098,9 +1098,80 @@ fn generate_triton_artifacts( let func_name = func_name.expect("Triton generator did not print FUNC_NAME"); let c_path = c_path.expect("Triton generator did not print C_PATH"); + patch_triton_aot_per_device_handles(&c_path); (func_name, c_path) } +fn patch_triton_aot_per_device_handles(c_path: &Path) { + let src = fs::read_to_string(c_path).expect("failed to read generated Triton C source"); + let module_name = src + .lines() + .find_map(|line| line.strip_prefix("CUmodule ")) + .and_then(|rest| rest.strip_suffix(" = NULL;")) + .map(str::to_string) + .expect("generated Triton C source should declare one CUmodule"); + let function_name = src + .lines() + .find_map(|line| line.strip_prefix("CUfunction ")) + .and_then(|rest| rest.strip_suffix(" = NULL;")) + .map(str::to_string) + .expect("generated Triton C source should declare one CUfunction"); + + let device_table_define = "#define OPENINFER_TRITON_DEVICE_TABLE_SIZE 16\n\n"; + let device_helper = "\n\nstatic CUresult openinfer_triton_current_device(CUdevice *dev) {\n CUresult err = cuCtxGetDevice(dev);\n if (err != CUDA_SUCCESS) {\n return err;\n }\n if (*dev < 0 || *dev >= OPENINFER_TRITON_DEVICE_TABLE_SIZE) {\n return CUDA_ERROR_INVALID_DEVICE;\n }\n return CUDA_SUCCESS;\n}"; + + let patched = src + .replace( + &format!("CUmodule {module_name} = NULL;"), + &format!( + "{device_table_define}CUmodule {module_name}[OPENINFER_TRITON_DEVICE_TABLE_SIZE] = {{0}};" + ), + ) + .replace( + &format!("CUfunction {function_name} = NULL;"), + &format!( + "CUfunction {function_name}[OPENINFER_TRITON_DEVICE_TABLE_SIZE] = {{0}};{device_helper}" + ), + ) + .replace( + &format!("CUDA_CHECK(cuModuleUnload({module_name}));"), + &format!( + "CUdevice dev = 0;\n CUDA_CHECK(openinfer_triton_current_device(&dev));\n if ({module_name}[dev] != NULL) {{\n CUDA_CHECK(cuModuleUnload({module_name}[dev]));\n {module_name}[dev] = NULL;\n {function_name}[dev] = NULL;\n }}" + ), + ); + let patched = patched + .replace( + &format!("cuModuleLoadData(&{module_name}, bin)"), + &format!("cuModuleLoadData(&{module_name}[dev], bin)"), + ) + .replace( + &format!("cuModuleGetFunction(&{function_name}, {module_name}, "), + &format!("cuModuleGetFunction(&{function_name}[dev], {module_name}[dev], "), + ) + .replace( + &format!("cuFuncSetCacheConfig({function_name}, "), + &format!("cuFuncSetCacheConfig({function_name}[dev], "), + ) + .replace( + &format!("cuFuncSetAttribute({function_name}, "), + &format!("cuFuncSetAttribute({function_name}[dev], "), + ) + .replace( + &format!("if ({function_name} == NULL)\n load_"), + &format!("CUdevice dev = 0;\n CUDA_CHECK(openinfer_triton_current_device(&dev));\n if ({function_name}[dev] == NULL)\n load_"), + ) + .replace( + &format!("return cuLaunchKernel({function_name}, "), + &format!("return cuLaunchKernel({function_name}[dev], "), + ) + .replace( + "int dev = 0;\n void *bin =", + "CUdevice dev = 0;\n CUDA_CHECK(openinfer_triton_current_device(&dev));\n void *bin =", + ); + + fs::write(c_path, patched).expect("failed to patch generated Triton C source"); +} + fn write_wrapper(generated_c: &Path, file_name: &str, wrapper_src: String) -> PathBuf { let wrapper_path = generated_c .parent() diff --git a/openinfer-qwen35-4b/Cargo.toml b/openinfer-qwen35-4b/Cargo.toml index 348d423f8..af1a8e4f0 100644 --- a/openinfer-qwen35-4b/Cargo.toml +++ b/openinfer-qwen35-4b/Cargo.toml @@ -21,9 +21,13 @@ tokio = { workspace = true, features = ["sync"] } [dev-dependencies] criterion = { workspace = true } +openinfer-vllm-frontend = { workspace = true } openinfer-vllm-support = { workspace = true } +reqwest = { workspace = true, features = ["json"] } sha2 = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["full"] } +tokio-util = { workspace = true } vllm-text = { workspace = true } [features] @@ -49,6 +53,10 @@ required-features = ["qwen35-4b"] name = "chunked_prefill" required-features = ["qwen35-4b"] +[[test]] +name = "serving_tp2" +required-features = ["qwen35-4b"] + [[bench]] name = "qwen35_ops" harness = false diff --git a/openinfer-qwen35-4b/src/batch_decode.rs b/openinfer-qwen35-4b/src/batch_decode.rs index 9611df01a..57f76c313 100644 --- a/openinfer-qwen35-4b/src/batch_decode.rs +++ b/openinfer-qwen35-4b/src/batch_decode.rs @@ -93,6 +93,9 @@ impl Qwen35Model { bufs: &mut BatchDecodeBuffers35, ) -> Result<()> { let eps = self.config.rms_norm_eps; + let tp = self.tensor_parallel; + let num_attention_heads = self.config.local_num_attention_heads(tp); + let num_key_value_heads = self.config.local_num_key_value_heads(tp); ops::gemm_into(&self.ctx, &attn.q_proj, &bufs.normed, &mut bufs.q_full); ops::gemm_into(&self.ctx, &attn.k_proj, &bufs.normed, &mut bufs.k_attn); @@ -108,8 +111,8 @@ impl Qwen35Model { &self.cos_cache, &self.sin_cache, &bufs.positions_d, - self.config.num_attention_heads, - self.config.num_key_value_heads, + num_attention_heads, + num_key_value_heads, self.config.rotary_dim, eps, ); @@ -130,7 +133,7 @@ impl Qwen35Model { &bufs.kv_tile_indices_d, &bufs.kv_chunk_size_d, &mut bufs.attn_out_full, - self.config.num_attention_heads, + num_attention_heads, bs, )?; @@ -140,7 +143,7 @@ impl Qwen35Model { crate::ffi::attention_gate_batch_hd256_cuda( qf_ptr as *const crate::ffi::Half, out_ptr as *mut crate::ffi::Half, - self.config.num_attention_heads as i32, + num_attention_heads as i32, bs as i32, self.ctx.stream.cu_stream(), ); @@ -152,6 +155,7 @@ impl Qwen35Model { &bufs.attn_out_full, &mut bufs.attn_results, ); + self.all_reduce_hidden(&mut bufs.attn_results)?; Ok(()) } @@ -220,9 +224,64 @@ impl Qwen35Model { &bufs.attn_out_full, &mut bufs.attn_results, ); + self.all_reduce_hidden(&mut bufs.attn_results)?; Ok(()) } + /// Eager batch decode step. + /// + /// Unlike `batch_decode_graph`, this does not pad to a CUDA Graph bucket and + /// does not capture/replay. Recurrent state is supplied directly by the + /// caller, which is the shape TP workers need for rank-local request state. + pub(crate) fn batch_decode_eager_logits( + &self, + token_ids: &[u32], + kv_states: &mut [&mut KvState], + recurrent_states: &mut [&mut RecurrentState], + bufs: &mut BatchDecodeBuffers35, + ) -> Result<()> { + let bs = token_ids.len(); + anyhow::ensure!( + bs > 0, + "batch_decode_eager_logits requires at least one request" + ); + anyhow::ensure!(bs == kv_states.len(), "token_ids / kv_states len mismatch"); + anyhow::ensure!( + bs == recurrent_states.len(), + "token_ids / recurrent_states len mismatch" + ); + anyhow::ensure!( + bs <= bufs.max_batch_size, + "batch size {bs} exceeds eager decode buffer capacity {}", + bufs.max_batch_size + ); + + let mut positions = Vec::with_capacity(bs); + for (i, kv) in kv_states.iter_mut().enumerate() { + let pos = kv.seq_len(); + self.ensure_rope_cache_covers(pos + 1)?; + kv.ensure_capacity(pos + 1)?; + kv.advance(1); + recurrent_states[i].seq_len += 1; + positions.push(pos as i32); + } + + bufs.set_batch_size(bs); + self.ctx + .stream + .memcpy_htod(token_ids, &mut bufs.token_ids_d)?; + self.ctx + .stream + .memcpy_htod(&positions, &mut bufs.positions_d)?; + + let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); + bufs.sync_paged_meta(&self.ctx, &kv_refs, bs)?; + + let kv_buffer = kv_states[0].buffer(); + let layout = *kv_states[0].layout(); + self.batch_decode_kernels_graph(kv_buffer, &layout, bs, recurrent_states, bufs) + } + // ========================================================================= // CUDA Graph batch decode // ========================================================================= @@ -306,11 +365,13 @@ impl Qwen35Model { // Take graphs out of graph_state to avoid split-borrow in the closure. let mut graphs = std::mem::take(&mut graph_state.graphs); let result = graphs[bucket_idx].run_or_capture(&self.ctx, || { + let mut slot_refs: Vec<&mut RecurrentState> = + graph_state.slot_states.iter_mut().collect(); self.batch_decode_kernels_graph( kv_buffer, &layout, padded_bs, - &mut graph_state.slot_states, + &mut slot_refs, &mut graph_state.buffers, ) }); @@ -397,8 +458,16 @@ impl Qwen35Model { self.config.num_key_value_heads, self.config.head_dim ); - let slot_states = &mut graph_state.slot_states; - self.batch_decode_batched_hybrid_kernels(kv_buffer, &layout, &plan, bs, slot_states, bufs) + let mut slot_states: Vec<&mut RecurrentState> = + graph_state.slot_states.iter_mut().collect(); + self.batch_decode_batched_hybrid_kernels( + kv_buffer, + &layout, + &plan, + bs, + &mut slot_states, + bufs, + ) } fn batch_decode_kernels_graph( @@ -406,7 +475,7 @@ impl Qwen35Model { kv_buffer: &cudarc::driver::CudaSlice, layout: &KvLayout, padded_bs: usize, - slot_states: &mut [RecurrentState], + slot_states: &mut [&mut RecurrentState], bufs: &mut BatchDecodeBuffers35, ) -> Result<()> { let eps = self.config.rms_norm_eps; @@ -476,6 +545,7 @@ impl Qwen35Model { &bufs.act_out, &mut bufs.mlp_out, ); + self.all_reduce_hidden(&mut bufs.mlp_out)?; ops::add_batch_into(&self.ctx, &bufs.hidden_mid, &bufs.mlp_out, &mut bufs.hidden)?; } @@ -506,7 +576,7 @@ impl Qwen35Model { layout: &KvLayout, plan: &ops::PrefillPagedPlan, bs: usize, - slot_states: &mut [RecurrentState], + slot_states: &mut [&mut RecurrentState], bufs: &mut BatchDecodeBuffers35, ) -> Result<()> { let eps = self.config.rms_norm_eps; @@ -621,7 +691,7 @@ impl Qwen35Model { fn batch_decode_linear_attention_slots( &self, attn: &LinearAttentionLayer, - slot_states: &mut [RecurrentState], + slot_states: &mut [&mut RecurrentState], layer_idx: usize, padded_bs: usize, bufs: &mut BatchDecodeBuffers35, @@ -632,6 +702,7 @@ impl Qwen35Model { ops::gemm_into(&self.ctx, &attn.in_proj_a, &bufs.normed, &mut bufs.a_proj); for (slot_idx, slot_state) in slot_states.iter_mut().enumerate().take(padded_bs) { + let slot_state = &mut **slot_state; let layer_state = &mut slot_state.layers[layer_idx]; ops::extract_vec_into(&self.ctx, &bufs.qkv, slot_idx, &mut bufs.qkv_tmp)?; diff --git a/openinfer-qwen35-4b/src/batch_decode_graph.rs b/openinfer-qwen35-4b/src/batch_decode_graph.rs index f6e795117..b415d6412 100644 --- a/openinfer-qwen35-4b/src/batch_decode_graph.rs +++ b/openinfer-qwen35-4b/src/batch_decode_graph.rs @@ -7,6 +7,7 @@ use openinfer_core::kv_pool::KvPool; use openinfer_core::tensor::DeviceContext; use super::config::Config35; +use super::config::TensorParallelConfig; use super::decode_buffers::BatchDecodeBuffers35; use super::recurrent_state::RecurrentState; @@ -59,14 +60,21 @@ impl BatchDecodeGraphState { pub(crate) fn with_capacity( ctx: &DeviceContext, config: &Config35, + tensor_parallel: TensorParallelConfig, kv_pool: &KvPool, max_batch: usize, ) -> Result { let padding_page_id = kv_pool.padding_page_id(); let max_total_pages = kv_pool.capacity_pages(); - let buffers = - BatchDecodeBuffers35::new(ctx, config, max_batch, max_total_pages, padding_page_id)?; + let buffers = BatchDecodeBuffers35::new( + ctx, + config, + tensor_parallel, + max_batch, + max_total_pages, + padding_page_id, + )?; let mut slot_states = Vec::with_capacity(max_batch); for _ in 0..max_batch { diff --git a/openinfer-qwen35-4b/src/config.rs b/openinfer-qwen35-4b/src/config.rs index 0b85dcd66..8f515f3c6 100644 --- a/openinfer-qwen35-4b/src/config.rs +++ b/openinfer-qwen35-4b/src/config.rs @@ -4,6 +4,21 @@ use serde::Deserialize; use std::collections::HashSet; use std::fs; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TensorParallelConfig { + pub(crate) rank: usize, + pub(crate) world_size: usize, +} + +impl Default for TensorParallelConfig { + fn default() -> Self { + Self { + rank: 0, + world_size: 1, + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum LayerType { FullAttention, @@ -197,11 +212,6 @@ impl Config35 { .count() } - /// Total Q dimension for full attention (includes gate). - pub(crate) fn full_attn_q_proj_dim(&self) -> usize { - self.num_attention_heads * self.head_dim * 2 - } - /// Q dimension for full attention (without gate). pub(crate) fn full_attn_q_dim(&self) -> usize { self.num_attention_heads * self.head_dim @@ -230,6 +240,218 @@ impl Config35 { pub(crate) fn linear_attn_z_dim(&self) -> usize { self.linear_num_value_heads * self.linear_value_head_dim } + + pub(crate) fn local_num_attention_heads(&self, tp: TensorParallelConfig) -> usize { + self.num_attention_heads / tp.world_size + } + + pub(crate) fn local_num_key_value_heads(&self, tp: TensorParallelConfig) -> usize { + self.num_key_value_heads / tp.world_size + } + + pub(crate) fn local_intermediate_size(&self, tp: TensorParallelConfig) -> usize { + self.intermediate_size / tp.world_size + } + + pub(crate) fn local_full_attn_q_dim(&self, tp: TensorParallelConfig) -> usize { + self.local_num_attention_heads(tp) * self.head_dim + } + + pub(crate) fn local_full_attn_kv_dim(&self, tp: TensorParallelConfig) -> usize { + self.local_num_key_value_heads(tp) * self.head_dim + } + + /// Local gated full-attention q projection output dimension. + pub(crate) fn local_full_attn_gated_q_dim(&self, tp: TensorParallelConfig) -> usize { + self.local_full_attn_q_dim(tp) * 2 + } +} + +impl TensorParallelConfig { + pub(crate) fn validate_for(self, config: &Config35, enable_cuda_graph: bool) -> Result<()> { + if self.world_size == 0 { + return Err(anyhow::anyhow!("tensor_parallel.world_size must be >= 1")); + } + if self.rank >= self.world_size { + return Err(anyhow::anyhow!( + "tensor_parallel.rank {} must be < world_size {}", + self.rank, + self.world_size + )); + } + if self.is_sharded() && enable_cuda_graph { + return Err(anyhow::anyhow!( + "Qwen3.5 tensor parallelism is eager-only in Phase 1; disable CUDA Graph for tp world_size={}", + self.world_size + )); + } + if !config.num_attention_heads.is_multiple_of(self.world_size) { + return Err(anyhow::anyhow!( + "num_attention_heads={} not divisible by tp world_size={}", + config.num_attention_heads, + self.world_size + )); + } + if !config.num_key_value_heads.is_multiple_of(self.world_size) { + return Err(anyhow::anyhow!( + "num_key_value_heads={} not divisible by tp world_size={}", + config.num_key_value_heads, + self.world_size + )); + } + if !config.intermediate_size.is_multiple_of(self.world_size) { + return Err(anyhow::anyhow!( + "intermediate_size={} not divisible by tp world_size={}", + config.intermediate_size, + self.world_size + )); + } + Ok(()) + } + + pub(crate) fn shard_range(self, total: usize) -> (usize, usize) { + let shard_len = total / self.world_size; + (self.rank * shard_len, shard_len) + } + + pub(crate) fn is_sharded(self) -> bool { + self.world_size > 1 + } +} + +#[cfg(test)] +mod tp_tests { + use super::*; + + fn test_config() -> Config35 { + Config35 { + hidden_size: 2560, + intermediate_size: 9216, + num_hidden_layers: 32, + vocab_size: 248320, + selection_vocab: 248320, + rms_norm_eps: 1e-6, + eos_token_id: 151645, + num_attention_heads: 16, + num_key_value_heads: 4, + head_dim: 256, + linear_num_key_heads: 16, + linear_key_head_dim: 128, + linear_num_value_heads: 32, + linear_value_head_dim: 128, + linear_conv_kernel_dim: 4, + rope_theta: 10_000.0, + rotary_dim: 64, + max_position_embeddings: 262_144, + tie_word_embeddings: true, + layer_types: vec![LayerType::LinearAttention; 32], + } + } + + #[test] + fn default_tensor_parallel_is_tp1() { + let config = test_config(); + let tp = TensorParallelConfig::default(); + + tp.validate_for(&config, true).unwrap(); + assert!(!tp.is_sharded()); + assert_eq!(tp.shard_range(config.full_attn_q_dim()), (0, 4096)); + assert_eq!(config.local_num_attention_heads(tp), 16); + assert_eq!(config.local_num_key_value_heads(tp), 4); + assert_eq!(config.local_intermediate_size(tp), 9216); + assert_eq!(config.local_full_attn_q_dim(tp), 4096); + assert_eq!(config.local_full_attn_kv_dim(tp), 1024); + assert_eq!(config.local_full_attn_gated_q_dim(tp), 8192); + } + + #[test] + fn computes_tp2_dense_local_dimensions() { + let config = test_config(); + let tp = TensorParallelConfig { + rank: 1, + world_size: 2, + }; + + tp.validate_for(&config, false).unwrap(); + assert!(tp.is_sharded()); + assert_eq!(tp.shard_range(config.full_attn_q_dim()), (2048, 2048)); + assert_eq!(config.local_num_attention_heads(tp), 8); + assert_eq!(config.local_num_key_value_heads(tp), 2); + assert_eq!(config.local_intermediate_size(tp), 4608); + assert_eq!(config.local_full_attn_q_dim(tp), 2048); + assert_eq!(config.local_full_attn_kv_dim(tp), 512); + assert_eq!(config.local_full_attn_gated_q_dim(tp), 4096); + } + + #[test] + fn rejects_invalid_world_size_and_rank() { + let config = test_config(); + + let err = TensorParallelConfig { + rank: 0, + world_size: 0, + } + .validate_for(&config, false) + .unwrap_err() + .to_string(); + assert!(err.contains("world_size must be >= 1")); + + let err = TensorParallelConfig { + rank: 2, + world_size: 2, + } + .validate_for(&config, false) + .unwrap_err() + .to_string(); + assert!(err.contains("rank 2 must be < world_size 2")); + } + + #[test] + fn rejects_indivisible_dense_dimensions() { + let tp = TensorParallelConfig { + rank: 0, + world_size: 3, + }; + + let mut config = test_config(); + let err = tp.validate_for(&config, false).unwrap_err().to_string(); + assert!(err.contains("num_attention_heads=16 not divisible")); + + config.num_attention_heads = 15; + config.num_key_value_heads = 4; + let err = tp.validate_for(&config, false).unwrap_err().to_string(); + assert!(err.contains("num_key_value_heads=4 not divisible")); + + config.num_key_value_heads = 3; + config.intermediate_size = 9217; + let err = tp.validate_for(&config, false).unwrap_err().to_string(); + assert!(err.contains("intermediate_size=9217 not divisible")); + } + + #[test] + fn rejects_tensor_parallel_cuda_graph_phase1() { + let config = test_config(); + let tp = TensorParallelConfig { + rank: 0, + world_size: 2, + }; + + let err = tp.validate_for(&config, true).unwrap_err().to_string(); + assert!(err.contains("eager-only in Phase 1")); + } + + #[test] + fn phase1_does_not_require_linear_attention_divisibility() { + let mut config = test_config(); + config.linear_num_key_heads = 17; + config.linear_num_value_heads = 31; + let tp = TensorParallelConfig { + rank: 1, + world_size: 2, + }; + + tp.validate_for(&config, false).unwrap(); + } } /// Schema kept identical to the pinned vLLM frontend; unread fields exist for diff --git a/openinfer-qwen35-4b/src/decode_buffers.rs b/openinfer-qwen35-4b/src/decode_buffers.rs index b8d3065ae..ef96f8c1a 100644 --- a/openinfer-qwen35-4b/src/decode_buffers.rs +++ b/openinfer-qwen35-4b/src/decode_buffers.rs @@ -4,7 +4,7 @@ use anyhow::Result; use cudarc::driver::CudaSlice; -use super::config::Config35; +use super::config::{Config35, TensorParallelConfig}; use openinfer_core::kv_pool::KvState; use openinfer_core::tensor::{DeviceContext, DeviceVec, HiddenStates}; @@ -70,19 +70,21 @@ impl BatchDecodeBuffers35 { pub(crate) fn new( ctx: &DeviceContext, config: &Config35, + tensor_parallel: TensorParallelConfig, max_batch_size: usize, max_total_pages: usize, padding_page_id: i32, ) -> Result { let h = config.hidden_size; let bs = max_batch_size; - let q_proj_dim = config.full_attn_q_proj_dim(); - let q_dim = config.full_attn_q_dim(); - let kv_dim = config.full_attn_kv_dim(); + let q_proj_dim = config.local_full_attn_gated_q_dim(tensor_parallel); + let q_dim = config.local_full_attn_q_dim(tensor_parallel); + let kv_dim = config.local_full_attn_kv_dim(tensor_parallel); let qkv_dim = config.linear_attn_qkv_dim(); let z_dim = config.linear_attn_z_dim(); let b_dim = config.linear_num_value_heads; let a_dim = b_dim; + let intermediate = config.local_intermediate_size(tensor_parallel); Ok(Self { max_batch_size: bs, @@ -90,8 +92,8 @@ impl BatchDecodeBuffers35 { normed: HiddenStates::zeros(ctx, h, bs)?, attn_results: HiddenStates::zeros(ctx, h, bs)?, hidden_mid: HiddenStates::zeros(ctx, h, bs)?, - gate_up_out: HiddenStates::zeros(ctx, 2 * config.intermediate_size, bs)?, - act_out: HiddenStates::zeros(ctx, config.intermediate_size, bs)?, + gate_up_out: HiddenStates::zeros(ctx, 2 * intermediate, bs)?, + act_out: HiddenStates::zeros(ctx, intermediate, bs)?, mlp_out: HiddenStates::zeros(ctx, h, bs)?, logits: HiddenStates::zeros(ctx, config.selection_vocab, bs)?, diff --git a/openinfer-qwen35-4b/src/executor.rs b/openinfer-qwen35-4b/src/executor.rs index cd64305f4..c0cb947be 100644 --- a/openinfer-qwen35-4b/src/executor.rs +++ b/openinfer-qwen35-4b/src/executor.rs @@ -33,9 +33,9 @@ impl RequestId { #[derive(Clone)] pub struct PrefillStepItem { - request_id: RequestId, - prompt_tokens: Vec, - logprobs: usize, + pub(crate) request_id: RequestId, + pub(crate) prompt_tokens: Vec, + pub(crate) logprobs: usize, } impl PrefillStepItem { @@ -50,9 +50,9 @@ impl PrefillStepItem { #[derive(Clone)] pub struct DecodeStepItem { - request_id: RequestId, - token_id: u32, - logprobs: usize, + pub(crate) request_id: RequestId, + pub(crate) token_id: u32, + pub(crate) logprobs: usize, } impl DecodeStepItem { @@ -89,10 +89,12 @@ pub struct DecodeRequestResult { pub logprob: Option, } +#[derive(Debug)] pub struct PrefillResult { pub requests: Vec, } +#[derive(Debug)] pub struct DecodeResult { pub requests: Vec, } diff --git a/openinfer-qwen35-4b/src/lib.rs b/openinfer-qwen35-4b/src/lib.rs index a9684b761..522916141 100644 --- a/openinfer-qwen35-4b/src/lib.rs +++ b/openinfer-qwen35-4b/src/lib.rs @@ -19,13 +19,14 @@ pub mod prefill_buffers; pub(crate) mod recurrent; pub(crate) mod recurrent_state; mod scheduler; +mod tp_executor; mod unified_forward; mod weights; use std::path::Path; use anyhow::{Result, anyhow}; -use openinfer_core::engine::{EngineHandle, EngineLoadOptions}; +use openinfer_core::engine::{EngineHandle, EngineLoadOptions, EpBackend}; pub use kernel_plan::kernel_plan; pub use scheduler::DEFAULT_MAX_PREFILL_TOKENS; @@ -40,6 +41,8 @@ pub mod runtime { DecodePlan, DecodeRequestResult, DecodeResult, DecodeStepItem, PrefillPlan, PrefillRequestResult, PrefillResult, PrefillStepItem, Qwen35Executor, RequestId, }; + pub use crate::scheduler::start_with_capacity; + pub use crate::tp_executor::Qwen35TpExecutor; pub use crate::weights::Qwen35Model; } @@ -50,12 +53,82 @@ pub mod runtime_ops { }; } -/// `max_batch` must be a decode bucket ({1,2,4,8,16,32,64}). pub fn start_engine( model_path: &Path, options: EngineLoadOptions, max_batch: usize, max_prefill_tokens: usize, +) -> Result { + start_engine_with_capacity(model_path, options, max_batch, max_prefill_tokens) +} + +#[derive(Clone, Debug)] +pub struct Qwen35LaunchOptions { + /// CUDA device for single-GPU loads (ignored when `tp_size > 1`). + pub device_ordinal: usize, + /// Tensor-parallel world size; `> 1` uses devices `0..tp_size`. + pub tp_size: usize, + /// TP Phase 1 supports eager-only multi-GPU execution. + pub cuda_graph: bool, + pub max_batch: usize, + pub max_prefill_tokens: usize, +} + +impl Qwen35LaunchOptions { + fn device_ordinals(&self) -> Result> { + anyhow::ensure!(self.tp_size >= 1, "Qwen3.5 tp_size must be >= 1"); + Ok(if self.tp_size == 1 { + vec![self.device_ordinal] + } else { + (0..self.tp_size).collect() + }) + } +} + +/// Start the Qwen3.5 engine for the server. TP Phase 1 supports eager-only +/// multi-GPU execution; single-GPU keeps the existing CUDA Graph-capable path. +pub fn launch( + model_path: &Path, + device_ordinal: usize, + cuda_graph: bool, + max_prefill_tokens: usize, +) -> Result { + launch_with_options( + model_path, + Qwen35LaunchOptions { + device_ordinal, + tp_size: 1, + cuda_graph, + max_batch: batch_decode_graph::MAX_BATCH, + max_prefill_tokens, + }, + ) +} + +pub fn launch_with_options( + model_path: &Path, + options: Qwen35LaunchOptions, +) -> Result { + let device_ordinals = options.device_ordinals()?; + start_engine_with_capacity( + model_path, + EngineLoadOptions { + enable_cuda_graph: options.cuda_graph, + device_ordinals, + parallel_config: None, + ep_backend: EpBackend::Nccl, + seed: 42, + }, + options.max_batch, + options.max_prefill_tokens, + ) +} + +pub fn start_engine_with_capacity( + model_path: &Path, + options: EngineLoadOptions, + max_batch: usize, + max_prefill_tokens: usize, ) -> Result { let EngineLoadOptions { enable_cuda_graph, @@ -63,6 +136,24 @@ pub fn start_engine( seed, .. } = options; + if device_ordinals.len() > 1 { + if enable_cuda_graph { + return Err(anyhow!( + "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" + )); + } + let model_path = model_path + .to_str() + .ok_or_else(|| anyhow!("model path must be valid UTF-8"))?; + return scheduler::start_tp_with_capacity( + model_path, + seed, + &device_ordinals, + max_batch, + max_prefill_tokens, + ); + } + anyhow::ensure!( enable_cuda_graph, "Qwen3.5 decode always captures CUDA Graphs; --cuda-graph=false is not supported" @@ -83,3 +174,22 @@ pub fn start_engine( let model = weights::Qwen35Model::from_safetensors(model_path, device_ordinal, max_batch)?; scheduler::start(model, seed, max_prefill_tokens) } + +#[cfg(test)] +mod tests { + use super::Qwen35LaunchOptions; + + #[test] + fn launch_options_reject_zero_tp_size() { + let options = Qwen35LaunchOptions { + device_ordinal: 3, + tp_size: 0, + cuda_graph: false, + max_batch: 1, + max_prefill_tokens: 1, + }; + + let err = options.device_ordinals().unwrap_err().to_string(); + assert!(err.contains("tp_size must be >= 1")); + } +} diff --git a/openinfer-qwen35-4b/src/prefill.rs b/openinfer-qwen35-4b/src/prefill.rs index 5c9a6318a..ddd30be90 100644 --- a/openinfer-qwen35-4b/src/prefill.rs +++ b/openinfer-qwen35-4b/src/prefill.rs @@ -170,13 +170,14 @@ impl Qwen35Model { kv_state.ensure_capacity(end_pos)?; kv_state.advance(seq_len); let kv_desc = kv_state.desc(); + let tp = self.tensor_parallel; let prefill_plan = PrefillPagedPlan::new( &self.ctx, &kv_desc, base_pos, seq_len, - c.num_attention_heads, - c.num_key_value_heads, + c.local_num_attention_heads(tp), + c.local_num_key_value_heads(tp), c.head_dim, )?; @@ -230,8 +231,9 @@ impl Qwen35Model { self.batched_rms_norm_offset(hidden_batch, &layer.input_layernorm, eps)?; // 2. Attention / Linear attention — per-token for correctness + let tp = self.tensor_parallel; let attn_out_dim = match &layer.attn { - LayerKind::FullAttention(_) => c.full_attn_q_dim(), + LayerKind::FullAttention(_) => c.local_full_attn_q_dim(tp), LayerKind::LinearAttention(_) => c.linear_attn_z_dim(), }; @@ -265,9 +267,10 @@ impl Qwen35Model { // 4. MLP (batched) let gate_up_out = ops::gemm(&self.ctx, &layer.mlp.gate_up_proj, &normed_batch)?; - let mut act_out = HiddenStates::zeros(&self.ctx, c.intermediate_size, seq_len)?; + let mut act_out = HiddenStates::zeros(&self.ctx, c.local_intermediate_size(tp), seq_len)?; ops::silu_mul_fused_batch_into(&self.ctx, &gate_up_out, &mut act_out)?; - let mlp_out = ops::gemm(&self.ctx, &layer.mlp.down_proj, &act_out)?; + let mut mlp_out = ops::gemm(&self.ctx, &layer.mlp.down_proj, &act_out)?; + self.all_reduce_hidden(&mut mlp_out)?; // 5. Residual ops::add_batch(&self.ctx, &hidden_plus_attn, &mlp_out) @@ -285,7 +288,10 @@ impl Qwen35Model { seq_len: usize, ) -> Result { let c = &self.config; - let attn_out_dim = c.full_attn_q_dim(); + let tp = self.tensor_parallel; + let num_attention_heads = c.local_num_attention_heads(tp); + let num_key_value_heads = c.local_num_key_value_heads(tp); + let attn_out_dim = c.local_full_attn_q_dim(tp); let eps = c.rms_norm_eps; let q_full_batch = ops::gemm(&self.ctx, &attn.q_proj, normed_batch)?; let k_batch = ops::gemm(&self.ctx, &attn.k_proj, normed_batch)?; @@ -332,8 +338,8 @@ impl Qwen35Model { layer_k_off, layer_v_off, pi_ptr as *const i32, - c.num_attention_heads as i32, - c.num_key_value_heads as i32, + num_attention_heads as i32, + num_key_value_heads as i32, seq_len as i32, sp_ptr as *const i32, c.rotary_dim as i32, @@ -381,8 +387,8 @@ impl Qwen35Model { kti_ptr as *const i32, kcs_ptr as *const i32, tnr_ptr as *const u32, - c.num_attention_heads as i32, - c.num_key_value_heads as i32, + num_attention_heads as i32, + num_key_value_heads as i32, HEAD_DIM as i32, layout.page_size as i32, seq_len as i32, @@ -408,7 +414,7 @@ impl Qwen35Model { ffi::attention_gate_batch_hd256_cuda( qf_ptr as *const ffi::Half, out_ptr as *mut ffi::Half, - c.num_attention_heads as i32, + num_attention_heads as i32, seq_len as i32, self.ctx.stream.cu_stream(), ); @@ -418,7 +424,9 @@ impl Qwen35Model { *full_idx += 1; // O projection (batched) - ops::gemm(&self.ctx, &attn.o_proj, &attn_out_batch) + let mut projected = ops::gemm(&self.ctx, &attn.o_proj, &attn_out_batch)?; + self.all_reduce_hidden(&mut projected)?; + Ok(projected) } fn prefill_linear_attention( diff --git a/openinfer-qwen35-4b/src/recurrent.rs b/openinfer-qwen35-4b/src/recurrent.rs index 5275eb754..1b44b066a 100644 --- a/openinfer-qwen35-4b/src/recurrent.rs +++ b/openinfer-qwen35-4b/src/recurrent.rs @@ -418,13 +418,15 @@ pub fn gated_delta_rule_prefill_chunkwise_into( &mut scratch.beta, num_key_heads, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill prepare failed: {e}"))?; gated_delta_rule_prefill_chunk_cumsum_inplace( ctx, &mut scratch.g_cumsum, qkv.seq_len, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill cumsum failed: {e}"))?; gated_delta_rule_prefill_chunk_a_into( ctx, &scratch.k_expanded, @@ -432,14 +434,16 @@ pub fn gated_delta_rule_prefill_chunkwise_into( &scratch.beta, &mut scratch.a_tril, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill A stage failed: {e}"))?; gated_delta_rule_prefill_chunk_solve_into( ctx, &scratch.a_tril, &mut scratch.a_inv, qkv.seq_len, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill solve failed: {e}"))?; gated_delta_rule_prefill_chunk_recompute_into( ctx, &scratch.k_expanded, @@ -450,7 +454,8 @@ pub fn gated_delta_rule_prefill_chunkwise_into( &scratch.a_inv, &scratch.g_cumsum, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill recompute failed: {e}"))?; gated_delta_rule_prefill_chunk_state_stage_into( ctx, &scratch.k_expanded, @@ -461,7 +466,8 @@ pub fn gated_delta_rule_prefill_chunkwise_into( &mut scratch.chunk_state, &mut scratch.v_new, num_value_heads, - )?; + ) + .map_err(|e| anyhow::anyhow!("GDR prefill state stage failed: {e}"))?; gated_delta_rule_prefill_chunk_o_stage_into( ctx, &scratch.q_expanded, @@ -473,6 +479,7 @@ pub fn gated_delta_rule_prefill_chunkwise_into( num_value_heads, 1.0 / (key_dim as f32).sqrt(), ) + .map_err(|e| anyhow::anyhow!("GDR prefill output stage failed: {e}")) } #[cfg(test)] diff --git a/openinfer-qwen35-4b/src/recurrent_state.rs b/openinfer-qwen35-4b/src/recurrent_state.rs index c5fe38c3b..0ebdfe38f 100644 --- a/openinfer-qwen35-4b/src/recurrent_state.rs +++ b/openinfer-qwen35-4b/src/recurrent_state.rs @@ -65,3 +65,20 @@ pub(crate) fn bytes_per_request(config: &Config35) -> usize { * (state_size * std::mem::size_of::() + conv_state_size * std::mem::size_of::()) } + +impl RecurrentState { + pub(crate) fn allocation_bytes(config: &Config35) -> usize { + bytes_per_request(config) + } +} + +#[cfg(test)] +mod tests { + #[test] + fn qwen35_4b_recurrent_allocation_is_49_125_mib() { + let bytes = 24 + * (32 * 128 * 128 * std::mem::size_of::() + + 8192 * 3 * std::mem::size_of::()); + assert_eq!(bytes, 49 * 1024 * 1024 + 128 * 1024); + } +} diff --git a/openinfer-qwen35-4b/src/scheduler.rs b/openinfer-qwen35-4b/src/scheduler.rs index 10d540951..be624fb9f 100644 --- a/openinfer-qwen35-4b/src/scheduler.rs +++ b/openinfer-qwen35-4b/src/scheduler.rs @@ -16,8 +16,12 @@ use rand::rngs::StdRng; use tokio::sync::mpsc; use crate::batch_decode_graph::BatchDecodeGraphState; +use crate::executor::{ + DecodeRequestResult, DecodeResult, PrefillRequestResult, PrefillResult, RequestId, +}; use crate::logprobs::snapshot_requested_logprobs; use crate::recurrent_state::RecurrentState; +use crate::tp_executor::{Qwen35TpExecutor, TpDecodeStepItem, TpPrefillChunkItem}; use crate::weights::Qwen35Model; use openinfer_core::engine::{ EngineHandle as SchedulerHandle, FinishReason, GenerateRequest as SchedulerRequest, KvCapacity, @@ -40,9 +44,7 @@ use self::plan::{ struct ActiveRequest35 { request_id: Option, token_tx: TokenSink, - kv: KvState, - /// Index into `BatchDecodeGraphState.slot_states`. - graph_slot_idx: usize, + backend_state: ActiveBackendState, last_token: u32, generated_count: usize, max_tokens: usize, @@ -57,14 +59,29 @@ struct ActiveRequest35 { /// at which point it is promoted into the decode batch. struct PrefillingRequest35 { req: SchedulerRequest, - kv: KvState, - rec: RecurrentState, + backend_state: PrefillBackendState, /// Prompt tokens prefilled so far. cursor: usize, /// Tokens to prefill in the step currently scheduled (set by `take_prefill_chunks`). step_chunk: usize, } +enum ActiveBackendState { + Single { + kv: KvState, + /// Index into `BatchDecodeGraphState.slot_states`. + graph_slot_idx: usize, + }, + Tp { + request_id: RequestId, + }, +} + +enum PrefillBackendState { + Single { kv: KvState, rec: RecurrentState }, + Tp { request_id: RequestId }, +} + pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; // ── Entry point ───────────────────────────────────────────────────────── @@ -73,6 +90,16 @@ pub(crate) fn start( model: Qwen35Model, seed: u64, max_prefill_tokens: usize, +) -> Result { + let max_batch = model.reserved_decode_slots; + start_with_capacity(model, seed, max_batch, max_prefill_tokens) +} + +pub fn start_with_capacity( + model: Qwen35Model, + seed: u64, + max_batch: usize, + max_prefill_tokens: usize, ) -> Result { assert!( max_prefill_tokens > 0, @@ -87,17 +114,22 @@ pub(crate) fn start( total_blocks, block_size, ); - let graph_state = model.create_batch_decode_graph_state()?; + let backend = SingleGpuBackend::new(model, max_batch)?; let (submit_tx, submit_rx) = mpsc::unbounded_channel(); let (startup_tx, startup_rx) = std_mpsc::channel(); let join_handle = thread::Builder::new() .name("scheduler-qwen35".into()) - .spawn(move || match bind_model_thread(&model) { + .spawn(move || match bind_model_thread(backend.model()) { Ok(_guard) => { let _ = startup_tx.send(Ok(())); - scheduler_loop(model, graph_state, submit_rx, seed, max_prefill_tokens); + scheduler_loop( + SchedulerBackend::Single(backend), + submit_rx, + seed, + max_prefill_tokens, + ); } Err(err) => { let _ = startup_tx.send(Err(err)); @@ -126,6 +158,559 @@ pub(crate) fn start( ) } +pub(crate) fn start_tp_with_capacity( + model_path: &str, + seed: u64, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, +) -> Result { + assert!( + max_prefill_tokens > 0, + "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" + ); + let backend = + TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; + let servable = servable_len( + backend.max_position_embeddings(), + backend.capacity_pages_for_requests(), + backend.page_size(), + ); + let kv_capacity = KvCapacity { + total_blocks: backend.capacity_pages_for_requests(), + block_size: backend.page_size(), + }; + + let (submit_tx, submit_rx) = mpsc::unbounded_channel(); + let join_handle = thread::Builder::new() + .name("scheduler-qwen35-tp".into()) + .spawn(move || { + scheduler_loop( + SchedulerBackend::Tp(backend), + submit_rx, + seed, + max_prefill_tokens, + ); + }) + .expect("failed to spawn Qwen3.5 TP scheduler thread"); + + Ok( + SchedulerHandle::new_with_join_handle(submit_tx, join_handle) + .with_servable_len(servable) + .with_kv_capacity(kv_capacity), + ) +} + +struct SingleGpuBackend { + model: Qwen35Model, + graph_state: BatchDecodeGraphState, +} + +enum SchedulerBackend { + Single(SingleGpuBackend), + Tp(TpSchedulerBackend), +} + +struct TpSchedulerBackend { + executor: Qwen35TpExecutor, + next_request_id: u64, +} + +impl SingleGpuBackend { + fn new(model: Qwen35Model, max_batch: usize) -> Result { + let graph_state = model.create_batch_decode_graph_state_with_capacity(max_batch)?; + Ok(Self { model, graph_state }) + } + + fn model(&self) -> &Qwen35Model { + &self.model + } + + fn max_batch(&self) -> usize { + self.graph_state.slot_states.len() + } + + fn page_size(&self) -> usize { + self.model.kv_pool().layout().page_size + } + + fn available_pages(&self) -> usize { + self.model.kv_pool().available_pages() + } + + fn capacity_pages_for_requests(&self) -> usize { + self.model.kv_pool().capacity_pages().saturating_sub(1) + } + + fn max_position_embeddings(&self) -> usize { + self.model.config().max_position_embeddings + } + + fn alloc_kv(&self) -> KvState { + self.model.alloc_kv() + } + + fn alloc_recurrent(&self) -> Result { + RecurrentState::new(self.model.device_ctx(), self.model.config()) + } + + fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(|w| w.as_slice()).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU prefill received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + self.model + .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) + } + + fn unified_step( + &mut self, + chunk: &mut ScheduledChunk, + active: &mut [ActiveRequest35], + ) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(|w| w.as_slice()).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU unified step received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); + let mut decode_kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU unified step received TP active state") + } + }) + .collect(); + self.model.unified_step( + &window_refs, + kvs, + &mut rec_refs, + &decode_tokens, + &mut decode_kv_refs, + &mut self.graph_state, + ) + } + + fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { + let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); + let mut kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU decode received TP active state") + } + }) + .collect(); + self.model + .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) + } + + fn sample_prefill_logits( + &mut self, + pending: &[SchedulerRequest], + logits: &HiddenStates, + rng: &mut StdRng, + ) -> Result<(Vec, Vec>)> { + debug_assert_eq!( + logits.seq_len, + pending.len(), + "Qwen3.5 prefill logits rows must preserve pending request order" + ); + let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); + let cpu_logits = + snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; + let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); + let sample_seed = rand::RngExt::random(rng); + let tokens = self.model.select_tokens_from_logits_varied( + logits, + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + openinfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + pending[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + fn sample_decode_logits( + &mut self, + active: &[ActiveRequest35], + rng: &mut StdRng, + ) -> Result<(Vec, Vec>)> { + let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); + let cpu_logits = snapshot_requested_logprobs( + self.model.device_ctx(), + &self.graph_state.buffers.logits, + &requested_logprobs, + )?; + let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); + let sample_seed = rand::RngExt::random(rng); + let tokens = self.model.select_tokens_batch_varied( + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + openinfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + active[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + fn is_stop_token(&self, token: u32) -> bool { + self.model.is_stop_token(token) + } + + fn copy_recurrent_to_slot( + &mut self, + recurrent: &RecurrentState, + slot_idx: usize, + ) -> Result<()> { + self.graph_state + .copy_state_to_slot(self.model.device_ctx(), recurrent, slot_idx) + } + + fn compact_slot(&mut self, active: &mut [ActiveRequest35], compaction: plan::SlotCompaction) { + let src_slot = match active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => graph_slot_idx, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + }; + debug_assert_eq!(src_slot, compaction.moved_from); + + let ctx = self.model.device_ctx(); + let src = &self.graph_state.slot_states[compaction.moved_from]; + for layer_idx in 0..src.layers.len() { + let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { + let (left, right) = self + .graph_state + .slot_states + .split_at_mut(compaction.moved_from); + ( + &right[0].layers[layer_idx], + &mut left[compaction.moved_to].layers[layer_idx], + ) + } else { + unreachable!("idx < active.len() <= last"); + }; + + ctx.stream + .memcpy_dtod(&src_part.state, &mut dst_part.state) + .expect("compact slot state copy failed"); + ctx.stream + .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) + .expect("compact slot conv_state copy failed"); + } + self.graph_state.slot_states[compaction.moved_to].seq_len = + self.graph_state.slot_states[compaction.moved_from].seq_len; + + match &mut active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => { + *graph_slot_idx = compaction.moved_to; + } + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + } + } +} + +impl TpSchedulerBackend { + fn new( + model_path: &str, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, + ) -> Result { + let executor = Qwen35TpExecutor::from_runtime_with_limits( + model_path, + false, + device_ordinals, + max_batch, + max_prefill_tokens, + )?; + Ok(Self { + executor, + next_request_id: 1, + }) + } + + fn alloc_request_id(&mut self) -> RequestId { + let id = RequestId::new(self.next_request_id); + self.next_request_id = self.next_request_id.wrapping_add(1).max(1); + id + } + + fn max_batch(&self) -> usize { + self.executor.max_batch() + } + + fn page_size(&self) -> usize { + self.executor.page_size() + } + + fn capacity_pages_for_requests(&self) -> usize { + self.executor.capacity_pages_for_requests() + } + + fn max_position_embeddings(&self) -> usize { + self.executor.max_position_embeddings() + } + + fn is_stop_token(&self, token: u32) -> bool { + self.executor.is_stop_token(token) + } + + fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + let page_size = self.page_size(); + let active_pages: usize = active + .iter() + .map(|req| pages_needed(current_active_tokens(req), page_size)) + .sum(); + let prefilling_pages: usize = prefilling + .iter() + .map(|req| pages_needed(req.cursor, page_size)) + .sum(); + self.capacity_pages_for_requests() + .saturating_sub(active_pages.saturating_add(prefilling_pages)) + } + + fn execute_prefill_chunk( + &self, + chunk: &ScheduledChunk, + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + anyhow::bail!("TP prefill received single-GPU chunk state"); + }; + let items: Vec = chunk + .reqs + .iter() + .zip(request_ids) + .zip(&chunk.windows) + .zip(&chunk.ends) + .map(|(((req, request_id), window), end)| { + TpPrefillChunkItem::new_with_sampling( + *request_id, + window.clone(), + req.logprobs, + req.params, + *end == req.prompt_tokens.len(), + ) + }) + .collect(); + let result = self + .executor + .execute_prefill_chunks_with_seed(&items, sample_seed)?; + align_prefill_results(chunk, &result) + } + + fn execute_decode( + &self, + active: &[ActiveRequest35], + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + let items: Vec = active + .iter() + .map(|req| { + let ActiveBackendState::Tp { request_id } = &req.backend_state else { + anyhow::bail!("TP decode received single-GPU active state"); + }; + Ok(TpDecodeStepItem::new( + *request_id, + req.last_token, + req.logprobs, + req.params, + )) + }) + .collect::>()?; + let result = self.executor.execute_decode_items(&items, sample_seed)?; + align_decode_results(active, &result) + } + + fn drop_request(&self, request_id: RequestId) { + if let Err(err) = self.executor.drop_request(request_id) { + warn!( + "failed to drop Qwen3.5 TP worker request {}: {err}", + request_id.get() + ); + } + } +} + +impl SchedulerBackend { + fn max_batch(&self) -> usize { + match self { + Self::Single(backend) => backend.max_batch(), + Self::Tp(backend) => backend.max_batch(), + } + } + + fn page_size(&self) -> usize { + match self { + Self::Single(backend) => backend.page_size(), + Self::Tp(backend) => backend.page_size(), + } + } + + fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + match self { + Self::Single(backend) => backend.available_pages(), + Self::Tp(backend) => backend.available_pages(active, prefilling), + } + } + + fn capacity_pages_for_requests(&self) -> usize { + match self { + Self::Single(backend) => backend.capacity_pages_for_requests(), + Self::Tp(backend) => backend.capacity_pages_for_requests(), + } + } + + fn max_position_embeddings(&self) -> usize { + match self { + Self::Single(backend) => backend.max_position_embeddings(), + Self::Tp(backend) => backend.max_position_embeddings(), + } + } + + fn alloc_prefill_state(&mut self) -> Result { + match self { + Self::Single(backend) => Ok(PrefillBackendState::Single { + kv: backend.alloc_kv(), + rec: backend.alloc_recurrent()?, + }), + Self::Tp(backend) => Ok(PrefillBackendState::Tp { + request_id: backend.alloc_request_id(), + }), + } + } + + fn is_tp(&self) -> bool { + matches!(self, Self::Tp(_)) + } + + fn is_stop_token(&self, token: u32) -> bool { + match self { + Self::Single(backend) => backend.is_stop_token(token), + Self::Tp(backend) => backend.is_stop_token(token), + } + } +} + +fn current_active_tokens(req: &ActiveRequest35) -> usize { + req.prompt_len + .saturating_add(req.generated_count.saturating_sub(1)) +} + +fn pages_needed(token_count: usize, page_size: usize) -> usize { + token_count.div_ceil(page_size) +} + +fn align_prefill_results( + chunk: &ScheduledChunk, + result: &PrefillResult, +) -> Result<(Vec, Vec>)> { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + anyhow::bail!("align_prefill_results requires TP chunk state"); + }; + let mut tokens = vec![0u32; chunk.reqs.len()]; + let mut logprobs = vec![None; chunk.reqs.len()]; + for PrefillRequestResult { + request_id, + first_token, + first_token_logprob, + } in &result.requests + { + let idx = request_ids + .iter() + .position(|id| id == request_id) + .ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP prefill returned unknown request id {}", + request_id.get() + ) + })?; + tokens[idx] = *first_token; + logprobs[idx] = first_token_logprob.clone(); + } + Ok((tokens, logprobs)) +} + +fn align_decode_results( + active: &[ActiveRequest35], + result: &DecodeResult, +) -> Result<(Vec, Vec>)> { + anyhow::ensure!( + active.len() == result.requests.len(), + "Qwen3.5 TP decode result row count mismatch: active={}, result={}", + active.len(), + result.requests.len() + ); + let mut tokens = Vec::with_capacity(active.len()); + let mut logprobs = Vec::with_capacity(active.len()); + for ( + active_req, + DecodeRequestResult { + request_id, + token, + logprob, + }, + ) in active.iter().zip(&result.requests) + { + let ActiveBackendState::Tp { + request_id: expected, + } = &active_req.backend_state + else { + anyhow::bail!("align_decode_results requires TP active state"); + }; + anyhow::ensure!( + *expected == *request_id, + "Qwen3.5 TP decode result request id mismatch: expected {}, got {}", + expected.get(), + request_id.get() + ); + tokens.push(*token); + logprobs.push(logprob.clone()); + } + Ok((tokens, logprobs)) +} + fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { max_context .min(max_pages.saturating_mul(page_size)) @@ -169,8 +754,7 @@ fn bind_model_thread(model: &Qwen35Model) -> Result { #[allow(clippy::needless_pass_by_value)] fn scheduler_loop( - model: Qwen35Model, - mut graph_state: BatchDecodeGraphState, + mut backend: SchedulerBackend, mut submit_rx: mpsc::UnboundedReceiver, seed: u64, prefill_budget: usize, @@ -179,7 +763,7 @@ fn scheduler_loop( let mut active: Vec = Vec::new(); let mut deferred: Vec = Vec::new(); let mut prefilling: Vec = Vec::new(); - let max_batch = graph_state.slot_states.len(); + let max_batch = backend.max_batch(); info!("scheduler ready (max_batch={})", max_batch); @@ -214,7 +798,7 @@ fn scheduler_loop( max_tokens: req.max_tokens, }) .collect(); - let page_size = model.kv_pool().layout().page_size; + let page_size = backend.page_size(); let prefilling_budget: Vec = prefilling .iter() .map(|p| PrefillKvBudget { @@ -223,9 +807,8 @@ fn scheduler_loop( max_tokens: p.req.max_tokens, }) .collect(); - let page_budget = model - .kv_pool() - .available_pages() + let page_budget = backend + .available_pages(&active, &prefilling) .saturating_sub(prefilling_future_pages(&prefilling_budget, page_size)); let decode_batching_slot = max_batch.saturating_sub(prefilling.len()); let admission = admit_pending_requests( @@ -236,8 +819,8 @@ fn scheduler_loop( page_budget, // KvPool capacity includes the CUDA Graph padding page reserved at // construction, so a real request can use at most the remaining pages. - model.kv_pool().capacity_pages().saturating_sub(1), - model.config().max_position_embeddings, + backend.capacity_pages_for_requests(), + backend.max_position_embeddings(), |req| req.prompt_tokens.len(), |req| req.max_tokens, ); @@ -253,10 +836,9 @@ fn scheduler_loop( req.prompt_tokens.len(), req.max_tokens ); - match RecurrentState::new(model.device_ctx(), model.config()) { - Ok(rec) => prefilling.push(PrefillingRequest35 { - kv: model.alloc_kv(), - rec, + match backend.alloc_prefill_state() { + Ok(backend_state) => prefilling.push(PrefillingRequest35 { + backend_state, cursor: 0, step_chunk: 0, req, @@ -277,32 +859,45 @@ fn scheduler_loop( // 5. Take this step's budgeted prefill chunk off the front of the queue, // then dispatch by plan. let scheduled = take_prefill_chunks(&mut prefilling, prefill_budget); - if let Some(plan) = plan::build_next_plan(!active.is_empty(), scheduled) { + let plan = if backend.is_tp() { + build_eager_only_plan(!active.is_empty(), scheduled) + } else { + plan::build_next_plan(!active.is_empty(), scheduled) + }; + if let Some(plan) = plan { match plan { ExecutionPlan::Unified { pending } => unified_step_sched( - &model, + &mut backend, &mut active, pending, &mut prefilling, - &mut graph_state, &mut rng, ), ExecutionPlan::Prefill { pending } => prefill_batch( - &model, + &mut backend, &mut active, pending, &mut prefilling, - &mut graph_state, &mut rng, ), ExecutionPlan::Decode => { - decode_step(&model, &mut active, &mut graph_state, &mut rng); + decode_step(&mut backend, &mut active, &mut rng); } } } } } +fn build_eager_only_plan(have_active: bool, pending: Vec) -> Option> { + if !pending.is_empty() { + Some(ExecutionPlan::Prefill { pending }) + } else if have_active { + Some(ExecutionPlan::Decode) + } else { + None + } +} + fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { let message = match reason { RejectReason::ContextLength { limit } => format!( @@ -329,117 +924,76 @@ fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { // ── Batch prefill ─────────────────────────────────────────────────────── fn prefill_batch( - model: &Qwen35Model, + backend: &mut SchedulerBackend, active: &mut Vec, scheduled: Vec, prefilling: &mut Vec, - graph_state: &mut BatchDecodeGraphState, rng: &mut StdRng, ) { let mut chunk = ScheduledChunk::from(scheduled); - // Scope the borrows of `chunk` to the executor call so the error path can - // move `chunk` into `fail_chunk`. - let result = { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let mut rec_refs: Vec<&mut RecurrentState> = chunk.recs.iter_mut().collect(); - model.batch_prefill_logits(&window_refs, &mut chunk.kvs, &mut rec_refs) - }; - let logits = match result { - Ok(v) => v, - Err(e) => { - warn!("batch prefill failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return; + let sample_seed = rand::RngExt::random(rng); + let (tokens, logprobs_vec) = match backend { + SchedulerBackend::Single(single) => { + // Scope the borrows of `chunk` to the executor call so the error path can + // move `chunk` into `fail_chunk`. + let logits = match single.batch_prefill_logits(&mut chunk) { + Ok(v) => v, + Err(e) => { + warn!("batch prefill failed: {e}"); + fail_chunk(chunk, &e.to_string()); + return; + } + }; + match single.sample_prefill_logits(&chunk.reqs, &logits, rng) { + Ok(v) => v, + Err(e) => { + warn!("prefill sampling failed: {e}"); + fail_chunk(chunk, &e.to_string()); + return; + } + } } - }; - - let (tokens, logprobs_vec) = - match sample_prefill_logits(model, &chunk.reqs, &logits, graph_state, rng) { + SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { Ok(v) => v, Err(e) => { - warn!("prefill sampling failed: {e}"); + warn!("TP prefill chunk failed: {e}"); + drop_tp_chunk_state(tp, &chunk); fail_chunk(chunk, &e.to_string()); return; } - }; - - promote_or_requeue( - model, - active, - prefilling, - graph_state, - chunk, - &tokens, - &logprobs_vec, - ); -} + }, + }; -fn sample_prefill_logits( - model: &Qwen35Model, - pending: &[SchedulerRequest], - logits: &HiddenStates, - graph_state: &mut BatchDecodeGraphState, - rng: &mut StdRng, -) -> Result<(Vec, Vec>)> { - debug_assert_eq!( - logits.seq_len, - pending.len(), - "Qwen3.5 prefill logits rows must preserve pending request order" - ); - let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); - let cpu_logits = snapshot_requested_logprobs(model.device_ctx(), logits, &requested_logprobs)?; - let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); - let sample_seed = rand::RngExt::random(rng); - let tokens = model.select_tokens_from_logits_varied( - logits, - &mut graph_state.buffers, - ¶ms_refs, - sample_seed, - )?; - - let logprobs = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - openinfer_sample::token_logprob_from_row( - &logits_f32, - tokens[i], - pending[i].logprobs, - ) - }) - }) - .collect(); - Ok((tokens, logprobs)) + promote_or_requeue(backend, active, prefilling, chunk, &tokens, &logprobs_vec); } // ── Unified step (prefill chunk + decode in one forward pass) ────────────── fn unified_step_sched( - model: &Qwen35Model, + backend: &mut SchedulerBackend, active: &mut Vec, scheduled: Vec, prefilling: &mut Vec, - graph_state: &mut BatchDecodeGraphState, rng: &mut StdRng, ) { + let SchedulerBackend::Single(backend) = backend else { + let chunk = ScheduledChunk::from(scheduled); + let message = "Qwen3.5 TP Phase 1 does not support unified prefill+decode steps"; + warn!("{message}"); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.to_string(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + fail_chunk(chunk, message); + return; + }; let mut chunk = ScheduledChunk::from(scheduled); // Scope the borrows of `chunk` / `active` to the executor call so the error // and decode-processing paths can use them afterwards. - let result = { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let mut rec_refs: Vec<&mut RecurrentState> = chunk.recs.iter_mut().collect(); - let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); - let mut decode_kv_refs: Vec<&mut KvState> = active.iter_mut().map(|r| &mut r.kv).collect(); - model.unified_step( - &window_refs, - &mut chunk.kvs, - &mut rec_refs, - &decode_tokens, - &mut decode_kv_refs, - graph_state, - ) - }; + let result = backend.unified_step(&mut chunk, active); let output = match result { Ok(v) => v, Err(e) => { @@ -460,7 +1014,7 @@ fn unified_step_sched( // Process decode results FIRST (it may retire requests and free graph slots // that promotion then fills densely). if output.decoded { - process_decode_logits(model, active, graph_state, rng); + process_decode_logits(backend, active, rng); } let prefill_logits = output @@ -468,7 +1022,7 @@ fn unified_step_sched( .as_ref() .expect("scheduled prefill chunk must return prefill logits"); let (tokens, logprobs_vec) = - match sample_prefill_logits(model, &chunk.reqs, prefill_logits, graph_state, rng) { + match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, rng) { Ok(v) => v, Err(e) => { warn!("unified prefill sampling failed: {e}"); @@ -477,73 +1031,58 @@ fn unified_step_sched( } }; - promote_or_requeue( - model, - active, - prefilling, - graph_state, - chunk, - &tokens, - &logprobs_vec, - ); + promote_or_requeue(backend, active, prefilling, chunk, &tokens, &logprobs_vec); } // ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── fn decode_step( - model: &Qwen35Model, + backend: &mut SchedulerBackend, active: &mut Vec, - graph_state: &mut BatchDecodeGraphState, rng: &mut StdRng, ) { - let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); - let mut kv_refs: Vec<&mut KvState> = active.iter_mut().map(|r| &mut r.kv).collect(); - - if let Err(e) = model.batch_decode_graph(&token_ids, &mut kv_refs, graph_state) { - warn!("batch_decode_graph error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return; - } - - // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) - let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); - let cpu_logits = match snapshot_requested_logprobs( - model.device_ctx(), - &graph_state.buffers.logits, - &requested_logprobs, - ) { - Ok(v) => v, - Err(e) => { - warn!("logprobs snapshot error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); + let sample_seed = rand::RngExt::random(rng); + let (tokens, logprobs_vec) = match backend { + SchedulerBackend::Single(single) => { + if let Err(e) = single.decode_graph(active) { + warn!("batch_decode_graph error: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + return; + } + // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) + match single.sample_decode_logits(active, rng) { + Ok(v) => v, + Err(e) => { + warn!("decode sampling/logprobs error: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + return; + } } - return; } - }; - - let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); - let sample_seed = rand::RngExt::random(rng); - let tokens = - match model.select_tokens_batch_varied(&mut graph_state.buffers, ¶ms_refs, sample_seed) - { - Ok(t) => t, + SchedulerBackend::Tp(tp) => match tp.execute_decode(active, sample_seed) { + Ok(v) => v, Err(e) => { - warn!("sampling error: {e}"); + warn!("TP eager decode error: {e}"); let message = e.to_string(); for req in active.drain(..) { + let state = req.backend_state; + if let ActiveBackendState::Tp { request_id } = state { + tp.drop_request(request_id); + } let _ = req.token_tx.send(TokenEvent::Error { message: message.clone(), prompt_tokens: req.prompt_len, @@ -552,37 +1091,22 @@ fn decode_step( } return; } - }; - - let logprobs_vec: Vec> = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - openinfer_sample::token_logprob_from_row(&logits_f32, tokens[i], active[i].logprobs) - }) - }) - .collect(); + }, + }; - dispatch_decode_tokens(model, active, &tokens, &logprobs_vec, graph_state); + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); } /// Process decode logits from unified step: sample, extract logprobs, dispatch. fn process_decode_logits( - model: &Qwen35Model, + backend: &mut SingleGpuBackend, active: &mut Vec, - graph_state: &mut BatchDecodeGraphState, rng: &mut StdRng, ) { - let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); - let cpu_logits = match snapshot_requested_logprobs( - model.device_ctx(), - &graph_state.buffers.logits, - &requested_logprobs, - ) { + let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, rng) { Ok(v) => v, Err(e) => { - warn!("decode logprobs snapshot error: {e}"); + warn!("decode sampling/logprobs error: {e}"); let message = e.to_string(); for req in active.drain(..) { let _ = req.token_tx.send(TokenEvent::Error { @@ -595,37 +1119,7 @@ fn process_decode_logits( } }; - let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); - let sample_seed = rand::RngExt::random(rng); - let tokens = - match model.select_tokens_batch_varied(&mut graph_state.buffers, ¶ms_refs, sample_seed) - { - Ok(t) => t, - Err(e) => { - warn!("decode sampling error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return; - } - }; - - let logprobs_vec: Vec> = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - openinfer_sample::token_logprob_from_row(&logits_f32, tokens[i], active[i].logprobs) - }) - }) - .collect(); - - dispatch_decode_tokens(model, active, &tokens, &logprobs_vec, graph_state); + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); } /// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. @@ -633,11 +1127,10 @@ fn process_decode_logits( /// `tokens` and `logprobs` are indexed by original position in `active`. /// Retirements collected first, then compacted in reverse order. fn dispatch_decode_tokens( - model: &Qwen35Model, + backend: &mut impl DecodeDispatchBackend, active: &mut Vec, tokens: &[u32], logprobs: &[Option], - graph_state: &mut BatchDecodeGraphState, ) { let n = active.len(); let mut to_retire = Vec::new(); @@ -648,7 +1141,7 @@ fn dispatch_decode_tokens( let req = &mut active[i]; req.generated_count += 1; - let is_eos = !req.params.ignore_eos && model.is_stop_token(token); + let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); let at_limit = req.generated_count >= req.max_tokens; if is_eos { @@ -697,57 +1190,58 @@ fn dispatch_decode_tokens( // Remove in reverse order so compact_slot indices stay valid for &i in to_retire.iter().rev() { - compact_slot(model, active, graph_state, i); + backend.retire_request(active, i); + } +} + +trait DecodeDispatchBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn retire_request(&mut self, active: &mut Vec, idx: usize); +} + +impl DecodeDispatchBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn retire_request(&mut self, active: &mut Vec, idx: usize) { + compact_single_slot(self, active, idx); + } +} + +impl DecodeDispatchBackend for SchedulerBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn retire_request(&mut self, active: &mut Vec, idx: usize) { + match self { + SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), + SchedulerBackend::Tp(backend) => { + let removed = active.swap_remove(idx); + if let ActiveBackendState::Tp { request_id } = removed.backend_state { + backend.drop_request(request_id); + } + } + } } } -/// Remove request at `idx` via swap_remove and compact graph slots. +/// Remove single-GPU request at `idx` via swap_remove and compact graph slots. /// /// After swap_remove, the element that was at `active.len()-1` (before remove) /// now sits at `idx`. Its graph slot must be copied into the vacated slot so /// that slots 0..active.len() remain dense. -fn compact_slot( - model: &Qwen35Model, +fn compact_single_slot( + backend: &mut SingleGpuBackend, active: &mut Vec, - graph_state: &mut BatchDecodeGraphState, idx: usize, ) { let compaction = compaction_after_retire(active.len(), idx); active.swap_remove(idx); if let Some(compaction) = compaction { - // The element that was at `last` is now at `idx`. - // Copy its recurrent state from slot `last` to slot `idx`. - let src_slot = active[idx].graph_slot_idx; - debug_assert_eq!(src_slot, compaction.moved_from); - - // D2D copy: graph_state.slot_states[src] -> graph_state.slot_states[dst] - // We can't borrow two slots mutably at once, so use raw index copy. - let ctx = model.device_ctx(); - let src = &graph_state.slot_states[compaction.moved_from]; - // Copy layer by layer using the public fields - for layer_idx in 0..src.layers.len() { - let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { - let (left, right) = graph_state.slot_states.split_at_mut(compaction.moved_from); - ( - &right[0].layers[layer_idx], - &mut left[compaction.moved_to].layers[layer_idx], - ) - } else { - unreachable!("idx < active.len() <= last"); - }; - - ctx.stream - .memcpy_dtod(&src_part.state, &mut dst_part.state) - .expect("compact slot state copy failed"); - ctx.stream - .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) - .expect("compact slot conv_state copy failed"); - } - graph_state.slot_states[compaction.moved_to].seq_len = - graph_state.slot_states[compaction.moved_from].seq_len; - - active[compaction.moved_to].graph_slot_idx = compaction.moved_to; + backend.compact_slot(active, compaction); } } @@ -756,21 +1250,41 @@ fn compact_slot( /// Step's scheduled prefill set struct ScheduledChunk { reqs: Vec, - kvs: Vec, - recs: Vec, + backend_state: ScheduledChunkBackendState, /// Prompt cursor after this step's chunk ends: Vec, /// This step's chunked token slice per request windows: Vec>, } +enum ScheduledChunkBackendState { + Single { + kvs: Vec, + recs: Vec, + }, + Tp { + request_ids: Vec, + }, +} + impl From> for ScheduledChunk { fn from(scheduled: Vec) -> Self { let n = scheduled.len(); + let is_tp = scheduled + .first() + .is_some_and(|p| matches!(p.backend_state, PrefillBackendState::Tp { .. })); let mut chunk = ScheduledChunk { reqs: Vec::with_capacity(n), - kvs: Vec::with_capacity(n), - recs: Vec::with_capacity(n), + backend_state: if is_tp { + ScheduledChunkBackendState::Tp { + request_ids: Vec::with_capacity(n), + } + } else { + ScheduledChunkBackendState::Single { + kvs: Vec::with_capacity(n), + recs: Vec::with_capacity(n), + } + }, ends: Vec::with_capacity(n), windows: Vec::with_capacity(n), }; @@ -781,8 +1295,20 @@ impl From> for ScheduledChunk { .push(p.req.prompt_tokens[p.cursor..end].to_vec()); chunk.ends.push(end); chunk.reqs.push(p.req); - chunk.kvs.push(p.kv); - chunk.recs.push(p.rec); + match (&mut chunk.backend_state, p.backend_state) { + ( + ScheduledChunkBackendState::Single { kvs, recs }, + PrefillBackendState::Single { kv, rec }, + ) => { + kvs.push(kv); + recs.push(rec); + } + ( + ScheduledChunkBackendState::Tp { request_ids }, + PrefillBackendState::Tp { request_id }, + ) => request_ids.push(request_id), + _ => unreachable!("mixed Qwen3.5 scheduler backend states in one chunk"), + } } chunk } @@ -817,35 +1343,44 @@ fn fail_chunk(chunk: ScheduledChunk, message: &str) { } } +fn drop_tp_chunk_state(backend: &TpSchedulerBackend, chunk: &ScheduledChunk) { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + return; + }; + for &request_id in request_ids { + backend.drop_request(request_id); + } +} + /// For each request in the just-prefilled chunk: if its prompt is now exhausted, /// sample its first token, emit events, and move it into the decode batch; /// otherwise re-queue it (with an advanced cursor) at the FRONT of `prefilling`. /// `tokens` / `logprobs` are indexed by request order in `chunk`. fn promote_or_requeue( - model: &Qwen35Model, + backend: &mut impl PrefillPromoteBackend, active: &mut Vec, prefilling: &mut Vec, - graph_state: &mut BatchDecodeGraphState, chunk: ScheduledChunk, tokens: &[u32], logprobs: &[Option], ) { let ScheduledChunk { reqs, - kvs, - recs, + backend_state, ends, .. } = chunk; let mut still_prefilling: Vec = Vec::new(); + let backend_states = split_scheduled_backend_state(backend_state); - for (i, (((req, kv), rec), end)) in reqs.into_iter().zip(kvs).zip(recs).zip(ends).enumerate() { + for (i, ((req, backend_state), end)) in + reqs.into_iter().zip(backend_states).zip(ends).enumerate() + { // Not finished: re-queue with the advanced cursor if end < req.prompt_tokens.len() { still_prefilling.push(PrefillingRequest35 { req, - kv, - rec, + backend_state, cursor: end, step_chunk: 0, }); @@ -864,7 +1399,7 @@ fn promote_or_requeue( }); } - if !req.params.ignore_eos && model.is_stop_token(first_token) { + if !req.params.ignore_eos && backend.is_stop_token(first_token) { debug!( "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", req.request_id, @@ -877,6 +1412,7 @@ fn promote_or_requeue( prompt_tokens: prompt_len, completion_tokens: 0, }); + backend.drop_prefill_state(backend_state); continue; } @@ -892,6 +1428,7 @@ fn promote_or_requeue( "request dropped: client disconnected: request_id={:?} tokens_generated={}", req.request_id, 0 ); + backend.drop_prefill_state(backend_state); continue; } @@ -908,20 +1445,15 @@ fn promote_or_requeue( prompt_tokens: prompt_len, completion_tokens: 1, }); + backend.drop_prefill_state(backend_state); continue; } - // Assign a graph slot and copy recurrent state into it. - let slot_idx = slot_for_new_request(active.len(), graph_state.slot_states.len()) - .expect("admission must reserve a graph slot"); - graph_state - .copy_state_to_slot(model.device_ctx(), &rec, slot_idx) - .expect("copy recurrent state to slot failed"); + let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); active.push(ActiveRequest35 { request_id: req.request_id, token_tx: req.token_tx, - kv, - graph_slot_idx: slot_idx, + backend_state: active_backend_state, last_token: first_token, generated_count: 1, max_tokens: req.max_tokens, @@ -934,5 +1466,95 @@ fn promote_or_requeue( prefilling.splice(0..0, still_prefilling); } +trait PrefillPromoteBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState; + fn drop_prefill_state(&mut self, state: PrefillBackendState); +} + +impl PrefillPromoteBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + let PrefillBackendState::Single { kv, rec } = state else { + panic!("single-GPU promotion received TP prefill state"); + }; + let slot_idx = slot_for_new_request(active_len, self.max_batch()) + .expect("admission must reserve a graph slot"); + self.copy_recurrent_to_slot(&rec, slot_idx) + .expect("copy recurrent state to slot failed"); + ActiveBackendState::Single { + kv, + graph_slot_idx: slot_idx, + } + } + + fn drop_prefill_state(&mut self, _state: PrefillBackendState) {} +} + +impl PrefillPromoteBackend for SchedulerBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + match (self, state) { + (SchedulerBackend::Single(single), PrefillBackendState::Single { kv, rec }) => { + let slot_idx = slot_for_new_request(active_len, single.max_batch()) + .expect("admission must reserve a graph slot"); + single + .copy_recurrent_to_slot(&rec, slot_idx) + .expect("copy recurrent state to slot failed"); + ActiveBackendState::Single { + kv, + graph_slot_idx: slot_idx, + } + } + (SchedulerBackend::Tp(_), PrefillBackendState::Tp { request_id }) => { + ActiveBackendState::Tp { request_id } + } + _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), + } + } + + fn drop_prefill_state(&mut self, state: PrefillBackendState) { + if let (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) = + (self, state) + { + backend.drop_request(request_id); + } + } +} + +fn split_scheduled_backend_state( + backend_state: ScheduledChunkBackendState, +) -> Vec { + match backend_state { + ScheduledChunkBackendState::Single { kvs, recs } => kvs + .into_iter() + .zip(recs) + .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) + .collect(), + ScheduledChunkBackendState::Tp { request_ids } => request_ids + .into_iter() + .map(|request_id| PrefillBackendState::Tp { request_id }) + .collect(), + } +} + #[cfg(test)] mod tests; diff --git a/openinfer-qwen35-4b/src/scheduler/tests.rs b/openinfer-qwen35-4b/src/scheduler/tests.rs index 1325d7412..679eb0259 100644 --- a/openinfer-qwen35-4b/src/scheduler/tests.rs +++ b/openinfer-qwen35-4b/src/scheduler/tests.rs @@ -1,4 +1,7 @@ use super::*; +use std::path::Path; + +use openinfer_core::engine::{EngineLoadOptions, EpBackend}; #[test] fn send_rejection_reports_kv_lifetime_request_tokens() { @@ -35,6 +38,104 @@ fn send_rejection_reports_kv_lifetime_request_tokens() { } } +#[test] +fn tp_scheduler_uses_eager_only_plan() { + let pending = vec!["prefill"]; + assert!( + matches!( + build_eager_only_plan(true, pending), + Some(ExecutionPlan::Prefill { pending }) if pending == vec!["prefill"] + ), + "TP Phase 1 should prefill first instead of choosing unified" + ); + assert!( + matches!( + build_eager_only_plan::<&str>(true, vec![]), + Some(ExecutionPlan::Decode) + ), + "TP Phase 1 should decode only when no prefill chunk is scheduled" + ); +} + +#[test] +fn tp_engine_rejects_cuda_graph_before_model_load() { + let err = match crate::start_engine_with_capacity( + Path::new("unused"), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0, 1], + parallel_config: None, + ep_backend: EpBackend::Nccl, + seed: 42, + }, + 1, + 1, + ) { + Ok(_) => panic!("TP CUDA Graph startup should fail"), + Err(err) => err.to_string(), + }; + assert!(err.contains("eager execution only")); +} + +#[test] +#[ignore = "requires two CUDA devices and Qwen3.5 weights"] +fn tp2_scheduler_chunked_prefill_then_decode_smoke() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let handle = + start_tp_with_capacity(&model_path, 42, &[0, 1], 1, 1).expect("start Qwen3.5 TP scheduler"); + let (token_tx, mut token_rx) = TokenSink::standalone(); + + handle + .submit(SchedulerRequest { + request_id: Some("tp2-scheduler-smoke".to_string()), + queued_at_unix_s: None, + data_parallel_rank: None, + prompt_tokens: vec![151_646, 9707], + params: SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens: 3, + lora_adapter: None, + token_tx, + logprobs: 1, + echo: false, + }) + .expect("submit TP scheduler request"); + + let mut tokens = Vec::new(); + loop { + match token_rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => { + let logprob = logprob.expect("TP scheduler smoke should return token logprob"); + assert!(logprob.logprob.is_finite()); + assert_eq!(logprob.top_logprobs.len(), 1); + tokens.push(id); + } + Some(TokenEvent::Finished { + finish_reason, + prompt_tokens, + completion_tokens, + }) => { + assert_eq!(finish_reason, FinishReason::Length); + assert_eq!(prompt_tokens, 2); + assert_eq!(completion_tokens, 3); + assert_eq!(tokens.len(), 3); + break; + } + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Error { message, .. }) => { + panic!("TP scheduler smoke failed: {message}") + } + Some(TokenEvent::Rejected { message, .. }) => { + panic!("TP scheduler smoke rejected: {message}") + } + None => panic!("TP scheduler channel closed before Finished"), + } + } +} + #[test] fn send_rejection_reports_context_window_limit() { let (token_tx, mut token_rx) = TokenSink::standalone(); diff --git a/openinfer-qwen35-4b/src/tp_executor.rs b/openinfer-qwen35-4b/src/tp_executor.rs new file mode 100644 index 000000000..2750db025 --- /dev/null +++ b/openinfer-qwen35-4b/src/tp_executor.rs @@ -0,0 +1,1616 @@ +//! Tensor-parallel worker runtime for Qwen3.5. +//! +//! Phase 1 supports eager dense TP prefill and decode. Unified execution still +//! fails closed until the scheduler path can drive ordered eager decode. + +use std::collections::HashSet; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex, mpsc}; +use std::thread::{self, JoinHandle}; + +use anyhow::Result; + +use crate::batch_decode_graph::MAX_BATCH; +use crate::config::TensorParallelConfig; +use crate::decode_buffers::BatchDecodeBuffers35; +use crate::executor::{ + DecodePlan, DecodeRequestResult, DecodeResult, DecodeStepItem, PrefillPlan, + PrefillRequestResult, PrefillResult, PrefillStepItem, RequestId, +}; +use crate::logprobs::snapshot_requested_logprobs; +use crate::prefill::PREFILL_CHUNK_LEN; +use crate::prefill_buffers::GdrChunkwiseScratch35; +use crate::recurrent_state::RecurrentState; +use crate::weights::{ModelRuntimeConfig, Qwen35Model}; +use openinfer_core::kv_pool::KvState; +use openinfer_core::sampler::SamplingParams; + +const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); +const TP_RUNTIME_STEP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); +const TP_WORKER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +const TP_RUNTIME_MEMORY_RESERVE_BYTES: usize = 512 * 1024 * 1024; + +#[allow(dead_code)] +enum TpWorkerCommand { + Ping { + resp: mpsc::Sender, + }, + RunPrefillChunks { + chunks: Vec, + sample_seed: u64, + resp: mpsc::Sender, + }, + RunDecodeStep { + requests: Vec, + sample_seed: u64, + resp: mpsc::Sender, + }, + RunUnifiedStep { + resp: mpsc::Sender, + }, + DropRequest { + request_id: RequestId, + resp: mpsc::Sender, + }, + Shutdown, +} + +#[derive(Debug)] +enum TpWorkerReply { + Ack, + Prefill(PrefillResult), + Decode(DecodeResult), +} + +#[derive(Debug)] +struct TpWorkerResponse { + rank: usize, + result: Result, +} + +#[derive(Default)] +struct TpRuntimePoison { + reason: Mutex>, +} + +impl TpRuntimePoison { + fn poison(&self, reason: String) -> String { + let mut current = self.reason.lock().unwrap_or_else(|err| err.into_inner()); + current.get_or_insert(reason).clone() + } + + fn ensure_healthy(&self) -> Result<()> { + if let Some(reason) = self + .reason + .lock() + .unwrap_or_else(|err| err.into_inner()) + .clone() + { + anyhow::bail!("Qwen3.5 TP executor is poisoned: {reason}"); + } + Ok(()) + } +} + +/// TP executor. Rank 0 is the primary worker and returns scheduler-visible +/// artifacts; every rank runs the same ordered state-mutating commands. +pub struct Qwen35TpExecutor { + workers: Vec, + poison: Arc, + world_size: usize, + max_batch: usize, + page_size: usize, + capacity_pages_for_requests: usize, + max_position_embeddings: usize, + eos_token_id: u32, +} + +#[derive(Clone)] +pub struct TpPrefillChunkItem { + request_id: RequestId, + prompt_tokens: Vec, + logprobs: usize, + sampling_params: SamplingParams, + finish_prefill: bool, +} + +impl TpPrefillChunkItem { + pub fn new( + request_id: RequestId, + prompt_tokens: Vec, + logprobs: usize, + finish_prefill: bool, + ) -> Self { + Self { + request_id, + prompt_tokens, + logprobs, + sampling_params: SamplingParams::default(), + finish_prefill, + } + } + + pub fn new_with_sampling( + request_id: RequestId, + prompt_tokens: Vec, + logprobs: usize, + sampling_params: SamplingParams, + finish_prefill: bool, + ) -> Self { + Self { + request_id, + prompt_tokens, + logprobs, + sampling_params, + finish_prefill, + } + } +} + +#[derive(Clone)] +pub struct TpDecodeStepItem { + request_id: RequestId, + token_id: u32, + logprobs: usize, + sampling_params: SamplingParams, +} + +impl TpDecodeStepItem { + pub fn new( + request_id: RequestId, + token_id: u32, + logprobs: usize, + sampling_params: SamplingParams, + ) -> Self { + Self { + request_id, + token_id, + logprobs, + sampling_params, + } + } +} + +impl Qwen35TpExecutor { + pub fn from_runtime( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + ) -> Result { + Self::from_runtime_with_capacity(model_path, enable_cuda_graph, device_ordinals, MAX_BATCH) + } + + pub fn from_runtime_with_capacity( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + max_batch: usize, + ) -> Result { + Self::from_runtime_with_limits( + model_path, + enable_cuda_graph, + device_ordinals, + max_batch, + PREFILL_CHUNK_LEN, + ) + } + + pub(crate) fn from_runtime_with_limits( + model_path: &str, + enable_cuda_graph: bool, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, + ) -> Result { + anyhow::ensure!( + device_ordinals.len() > 1, + "Qwen3.5 TP executor requires at least two CUDA devices, got {}", + device_ordinals.len() + ); + anyhow::ensure!( + !enable_cuda_graph, + "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" + ); + anyhow::ensure!( + max_prefill_tokens > 0, + "Qwen3.5 TP max_prefill_tokens must be positive" + ); + + let world_size = device_ordinals.len(); + let mut models = Vec::with_capacity(world_size); + for (rank, &device_ordinal) in device_ordinals.iter().enumerate() { + models.push(Qwen35Model::from_safetensors_with_runtime( + model_path, + ModelRuntimeConfig { + enable_cuda_graph: false, + tensor_parallel: Some(TensorParallelConfig { rank, world_size }), + device_ordinal, + }, + )?); + } + let first = models + .first() + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP executor loaded no models"))?; + let page_size = first.kv_pool().layout().page_size; + let mut min_capacity_pages = usize::MAX; + for (rank, model) in models.iter().enumerate() { + let rank_page_size = model.kv_pool().layout().page_size; + anyhow::ensure!( + rank_page_size == page_size, + "Qwen3.5 TP rank {rank} KV page size {rank_page_size} does not match rank 0 page size {page_size}" + ); + min_capacity_pages = min_capacity_pages.min(model.kv_pool().capacity_pages()); + } + let capacity_pages_for_requests = min_capacity_pages.saturating_sub(1); + let max_position_embeddings = first.config().max_position_embeddings; + let eos_token_id = first.config().eos_token_id; + + let nccl_id = cudarc::nccl::safe::Id::new() + .map_err(|e| anyhow::anyhow!("failed to create Qwen3.5 TP NCCL id: {e:?}"))?; + let startup_gate = Arc::new(TpStartupGate::default()); + let effective_max_batch = Arc::new(AtomicUsize::new(0)); + let poison = Arc::new(TpRuntimePoison::default()); + let mut workers = Vec::with_capacity(world_size); + let mut preflights = Vec::with_capacity(world_size); + let mut startups = Vec::with_capacity(world_size); + for (rank, model) in models.into_iter().enumerate() { + match TpWorker::spawn( + rank, + world_size, + model, + max_batch, + max_prefill_tokens, + nccl_id, + Arc::clone(&startup_gate), + Arc::clone(&effective_max_batch), + Arc::clone(&poison), + ) { + Ok((worker, preflight, startup)) => { + workers.push(worker); + preflights.push(preflight); + startups.push(startup); + } + Err(err) => { + startup_gate.cancel(); + return Err(err); + } + } + } + let mut min_rank_max_batch = max_batch; + for (rank, preflight) in preflights.into_iter().enumerate() { + match preflight.recv() { + Ok(Ok(rank_max_batch)) => { + min_rank_max_batch = min_rank_max_batch.min(rank_max_batch); + } + Ok(Err(err)) => { + startup_gate.cancel(); + return Err(err); + } + Err(_) => { + startup_gate.cancel(); + return Err(anyhow::anyhow!( + "Qwen3.5 TP worker {rank} exited during pre-NCCL startup" + )); + } + } + } + anyhow::ensure!( + min_rank_max_batch > 0, + "Qwen3.5 TP has no memory capacity for one recurrent request state" + ); + effective_max_batch.store(min_rank_max_batch, Ordering::Release); + if min_rank_max_batch < max_batch { + log::warn!( + "Qwen3.5 TP max_batch reduced from {max_batch} to {min_rank_max_batch} by rank-local recurrent-state memory capacity" + ); + } + let (watchdog_done, watchdog) = match spawn_nccl_startup_watchdog() { + Ok(watchdog) => watchdog, + Err(err) => { + startup_gate.cancel(); + return Err(err); + } + }; + startup_gate.connect(); + let startup_result = startups + .into_iter() + .enumerate() + .try_for_each(|(rank, startup)| { + startup.recv().map_err(|_| { + anyhow::anyhow!("Qwen3.5 TP worker {rank} exited during startup") + })? + }); + if let Err(err) = startup_result { + drop(workers); + disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; + return Err(err); + } + disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; + + Ok(Self { + workers, + poison, + world_size, + max_batch: min_rank_max_batch, + page_size, + capacity_pages_for_requests, + max_position_embeddings, + eos_token_id, + }) + } + + pub fn world_size(&self) -> usize { + self.world_size + } + + pub fn max_batch(&self) -> usize { + self.max_batch + } + + pub fn page_size(&self) -> usize { + self.page_size + } + + pub fn capacity_pages_for_requests(&self) -> usize { + self.capacity_pages_for_requests + } + + pub fn max_position_embeddings(&self) -> usize { + self.max_position_embeddings + } + + pub fn is_stop_token(&self, token_id: u32) -> bool { + token_id == self.eos_token_id + } + + pub fn ping_all(&self) -> Result<()> { + self.poison.ensure_healthy()?; + self.broadcast_ack(TpWorkerCommandKind::Ping) + } + + pub fn execute_prefill(&self, plan: PrefillPlan<'_>) -> Result { + anyhow::ensure!( + !plan.requests.is_empty(), + "Qwen3.5 TP prefill plan requires at least one request" + ); + let chunks: Vec = plan + .requests + .iter() + .cloned() + .map(TpPrefillChunkItem::from) + .collect(); + self.execute_prefill_chunks(&chunks) + } + + pub fn execute_prefill_chunks(&self, chunks: &[TpPrefillChunkItem]) -> Result { + self.execute_prefill_chunks_with_seed(chunks, 0) + } + + pub fn execute_prefill_chunks_with_seed( + &self, + chunks: &[TpPrefillChunkItem], + sample_seed: u64, + ) -> Result { + self.poison.ensure_healthy()?; + anyhow::ensure!( + !chunks.is_empty(), + "Qwen3.5 TP prefill chunk command requires at least one chunk" + ); + let chunks = chunks.to_vec(); + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + self.send_or_poison( + worker, + TpWorkerCommand::RunPrefillChunks { + chunks: chunks.clone(), + sample_seed, + resp: resp_tx.clone(), + }, + )?; + } + drop(resp_tx); + wait_for_prefill(resp_rx, self.workers.len(), &self.poison) + } + + pub fn execute_decode(&self, plan: DecodePlan<'_>) -> Result { + anyhow::ensure!( + !plan.requests.is_empty(), + "Qwen3.5 TP decode plan requires at least one request" + ); + let requests: Vec = plan + .requests + .iter() + .map(|request| { + TpDecodeStepItem::new( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + ) + }) + .collect(); + self.execute_decode_items(&requests, 0) + } + + pub fn execute_decode_items( + &self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result { + self.poison.ensure_healthy()?; + anyhow::ensure!( + !requests.is_empty(), + "Qwen3.5 TP decode plan requires at least one request" + ); + let requests = requests.to_vec(); + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + self.send_or_poison( + worker, + TpWorkerCommand::RunDecodeStep { + requests: requests.clone(), + sample_seed, + resp: resp_tx.clone(), + }, + )?; + } + drop(resp_tx); + wait_for_decode(resp_rx, self.workers.len(), &self.poison) + } + + pub fn drop_request(&self, request_id: RequestId) -> Result<()> { + self.poison.ensure_healthy()?; + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + self.send_or_poison( + worker, + TpWorkerCommand::DropRequest { + request_id, + resp: resp_tx.clone(), + }, + )?; + } + drop(resp_tx); + wait_for_acks(resp_rx, self.workers.len(), "drop request", &self.poison) + } + + fn broadcast_ack(&self, kind: TpWorkerCommandKind) -> Result<()> { + let (resp_tx, resp_rx) = mpsc::channel(); + for worker in &self.workers { + let command = match kind { + TpWorkerCommandKind::Ping => TpWorkerCommand::Ping { + resp: resp_tx.clone(), + }, + TpWorkerCommandKind::RunPrefillChunks => TpWorkerCommand::RunPrefillChunks { + chunks: Vec::new(), + sample_seed: 0, + resp: resp_tx.clone(), + }, + TpWorkerCommandKind::RunDecodeStep => TpWorkerCommand::RunDecodeStep { + requests: Vec::new(), + sample_seed: 0, + resp: resp_tx.clone(), + }, + TpWorkerCommandKind::RunUnifiedStep => TpWorkerCommand::RunUnifiedStep { + resp: resp_tx.clone(), + }, + }; + self.send_or_poison(worker, command)?; + } + drop(resp_tx); + wait_for_acks(resp_rx, self.workers.len(), kind.name(), &self.poison) + } + + fn send_or_poison(&self, worker: &TpWorker, command: TpWorkerCommand) -> Result<()> { + worker.send(command).map_err(|err| { + let reason = self + .poison + .poison(format!("failed to dispatch TP worker command: {err:#}")); + anyhow::anyhow!(reason) + }) + } +} + +impl Drop for Qwen35TpExecutor { + fn drop(&mut self) { + for worker in &self.workers { + let _ = worker.tx.send(TpWorkerCommand::Shutdown); + } + for worker in &mut self.workers { + worker.join_bounded(); + } + } +} + +#[allow(dead_code)] +#[derive(Clone, Copy)] +enum TpWorkerCommandKind { + Ping, + RunPrefillChunks, + RunDecodeStep, + RunUnifiedStep, +} + +impl TpWorkerCommandKind { + fn name(self) -> &'static str { + match self { + Self::Ping => "ping", + Self::RunPrefillChunks => "prefill chunks", + Self::RunDecodeStep => "decode step", + Self::RunUnifiedStep => "unified step", + } + } +} + +fn spawn_nccl_startup_watchdog() -> Result<(mpsc::SyncSender<()>, JoinHandle<()>)> { + let (done_tx, done_rx) = mpsc::sync_channel(1); + let watchdog = thread::Builder::new() + .name("qwen35-tp-nccl-startup-watchdog".into()) + .spawn(move || { + if done_rx.recv_timeout(TP_NCCL_STARTUP_TIMEOUT).is_ok() { + return; + } + eprintln!( + "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", + TP_NCCL_STARTUP_TIMEOUT.as_secs() + ); + log::error!( + "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", + TP_NCCL_STARTUP_TIMEOUT.as_secs() + ); + std::process::abort(); + }) + .map_err(|err| anyhow::anyhow!("failed to spawn Qwen3.5 TP NCCL watchdog: {err}"))?; + Ok((done_tx, watchdog)) +} + +fn disarm_nccl_startup_watchdog( + done_tx: mpsc::SyncSender<()>, + watchdog: JoinHandle<()>, +) -> Result<()> { + done_tx + .send(()) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog exited unexpectedly"))?; + watchdog + .join() + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog panicked")) +} + +struct TpWorker { + tx: mpsc::Sender, + handle: Option>, + done: mpsc::Receiver<()>, +} + +#[derive(Clone, Copy, Default, PartialEq, Eq)] +enum TpStartupDecision { + #[default] + Pending, + Connect, + Cancel, +} + +#[derive(Default)] +struct TpStartupGate { + decision: Mutex, + changed: Condvar, +} + +impl TpStartupGate { + fn connect(&self) { + self.set(TpStartupDecision::Connect); + } + + fn cancel(&self) { + self.set(TpStartupDecision::Cancel); + } + + fn wait(&self) -> bool { + let mut decision = self.decision.lock().unwrap_or_else(|err| err.into_inner()); + while *decision == TpStartupDecision::Pending { + decision = self + .changed + .wait(decision) + .unwrap_or_else(|err| err.into_inner()); + } + *decision == TpStartupDecision::Connect + } + + fn set(&self, next: TpStartupDecision) { + let mut decision = self.decision.lock().unwrap_or_else(|err| err.into_inner()); + if *decision == TpStartupDecision::Pending { + *decision = next; + self.changed.notify_all(); + } + } +} + +impl TpWorker { + fn spawn( + rank: usize, + world_size: usize, + model: Qwen35Model, + max_batch: usize, + max_prefill_tokens: usize, + nccl_id: cudarc::nccl::safe::Id, + startup_gate: Arc, + effective_max_batch: Arc, + poison: Arc, + ) -> Result<( + Self, + mpsc::Receiver>, + mpsc::Receiver>, + )> { + let (tx, rx) = mpsc::channel(); + let (preflight_tx, preflight_rx) = mpsc::channel(); + let (startup_tx, startup_rx) = mpsc::channel(); + let (done_tx, done_rx) = mpsc::channel(); + let panic_poison = Arc::clone(&poison); + let handle = thread::Builder::new() + .name(format!("qwen35-tp-rank-{rank}")) + .spawn(move || { + let outcome = catch_unwind(AssertUnwindSafe(|| { + let prepared = TpWorkerPrepared::new( + rank, + world_size, + model, + max_batch, + max_prefill_tokens, + ); + let prepared = match prepared { + Ok((prepared, rank_max_batch)) => { + let _ = preflight_tx.send(Ok(rank_max_batch)); + prepared + } + Err(err) => { + let _ = preflight_tx.send(Err(err)); + return; + } + }; + if !startup_gate.wait() { + return; + } + let max_batch = effective_max_batch.load(Ordering::Acquire); + match prepared.connect(nccl_id, max_batch, poison) { + Ok(mut state) => { + let _ = startup_tx.send(Ok(())); + state.run(rx); + } + Err(err) => { + let _ = startup_tx.send(Err(err)); + } + } + })); + if outcome.is_err() { + panic_poison.poison(format!("worker rank {rank} panicked")); + } + let _ = done_tx.send(()); + }) + .map_err(|e| anyhow::anyhow!("failed to spawn Qwen3.5 TP worker {rank}: {e}"))?; + + Ok(( + Self { + tx, + handle: Some(handle), + done: done_rx, + }, + preflight_rx, + startup_rx, + )) + } + + fn send(&self, command: TpWorkerCommand) -> Result<()> { + self.tx + .send(command) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP worker channel closed")) + } + + fn join_bounded(&mut self) { + if self.handle.is_none() { + return; + } + if self.done.recv_timeout(TP_WORKER_SHUTDOWN_TIMEOUT).is_err() { + fatal_tp_abort("Qwen3.5 TP worker did not exit during bounded shutdown"); + } + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } +} + +impl Drop for TpWorker { + fn drop(&mut self) { + let _ = self.tx.send(TpWorkerCommand::Shutdown); + self.join_bounded(); + } +} + +struct TpWorkerState { + rank: usize, + _world_size: usize, + max_batch: usize, + model: Qwen35Model, + requests: Vec, + decode_buffers: BatchDecodeBuffers35, + sample_scratch: openinfer_sample::SampleScratch, + _cublas_guard: CublasThreadGuard, + poison: Arc, +} + +struct TpWorkerPrepared { + rank: usize, + world_size: usize, + max_batch: usize, + model: Qwen35Model, + decode_buffers: BatchDecodeBuffers35, + sample_scratch: openinfer_sample::SampleScratch, + cublas_guard: CublasThreadGuard, +} + +struct TpRequestState { + request_id: RequestId, + phase: TpRequestPhase, + kv: KvState, + recurrent: RecurrentState, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TpRequestPhase { + Prefilling, + Decoding, +} + +impl TpWorkerPrepared { + fn new( + rank: usize, + world_size: usize, + model: Qwen35Model, + requested_max_batch: usize, + max_prefill_tokens: usize, + ) -> Result<(Self, usize)> { + let cublas_guard = bind_worker_thread(&model)?; + let (free_bytes, total_bytes) = model + .device_ctx() + .ctx + .mem_get_info() + .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; + let recurrent_bytes = RecurrentState::allocation_bytes(model.config()); + let prefill_scratch_tokens = prefill_scratch_tokens(max_prefill_tokens); + let prefill_scratch_bytes = + GdrChunkwiseScratch35::estimate_bytes(model.config(), prefill_scratch_tokens); + let max_batch = effective_recurrent_capacity( + requested_max_batch, + free_bytes, + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ); + anyhow::ensure!( + max_batch > 0, + "Qwen3.5 TP rank {rank} has {} MiB free after fixed buffers, but one recurrent request needs {} MiB plus {} MiB runtime reserve and {} MiB prefill scratch for {} tokens", + free_bytes / (1024 * 1024), + recurrent_bytes / (1024 * 1024), + TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), + prefill_scratch_bytes / (1024 * 1024), + prefill_scratch_tokens, + ); + log::info!( + "Qwen3.5 TP rank {rank} recurrent capacity: requested={requested_max_batch}, effective={max_batch}, per_request={:.3} MiB, free={:.0} MiB/{:.0} MiB, runtime_reserve={} MiB, prefill_tokens={}, prefill_scratch={:.0} MiB", + recurrent_bytes as f64 / 1024.0 / 1024.0, + free_bytes as f64 / 1024.0 / 1024.0, + total_bytes as f64 / 1024.0 / 1024.0, + TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), + prefill_scratch_tokens, + prefill_scratch_bytes as f64 / 1024.0 / 1024.0, + ); + let decode_buffers = model.create_batch_decode_buffers_with_capacity(max_batch)?; + let sample_scratch = openinfer_sample::SampleScratch::new( + model.device_ctx(), + model.config().vocab_size, + max_batch, + )?; + Ok(( + Self { + rank, + world_size, + max_batch, + model, + decode_buffers, + sample_scratch, + cublas_guard, + }, + max_batch, + )) + } + + fn connect( + self, + nccl_id: cudarc::nccl::safe::Id, + effective_max_batch: usize, + poison: Arc, + ) -> Result { + let Self { + rank, + world_size, + max_batch, + mut model, + decode_buffers, + sample_scratch, + cublas_guard, + } = self; + anyhow::ensure!( + effective_max_batch > 0 && effective_max_batch <= max_batch, + "Qwen3.5 TP rank {rank} effective max_batch {effective_max_batch} exceeds local capacity {max_batch}" + ); + let comm = cudarc::nccl::safe::Comm::from_rank( + model.device_ctx().stream.clone(), + rank, + world_size, + nccl_id, + ) + .map_err(|e| anyhow::anyhow!("failed to initialize Qwen3.5 TP NCCL rank {rank}: {e:?}"))?; + model.attach_tp_comm(comm); + Ok(TpWorkerState { + rank, + _world_size: world_size, + max_batch: effective_max_batch, + model, + requests: Vec::new(), + decode_buffers, + sample_scratch, + _cublas_guard: cublas_guard, + poison, + }) + } +} + +fn prefill_scratch_tokens(max_prefill_tokens: usize) -> usize { + max_prefill_tokens.min(PREFILL_CHUNK_LEN) +} + +fn effective_recurrent_capacity( + requested_max_batch: usize, + free_bytes: usize, + recurrent_bytes_per_request: usize, + runtime_reserve_bytes: usize, + prefill_scratch_bytes: usize, +) -> usize { + if recurrent_bytes_per_request == 0 { + return requested_max_batch; + } + requested_max_batch.min( + free_bytes + .saturating_sub(runtime_reserve_bytes) + .saturating_sub(prefill_scratch_bytes) + / recurrent_bytes_per_request, + ) +} + +impl TpWorkerState { + fn run(&mut self, rx: mpsc::Receiver) { + while let Ok(command) = rx.recv() { + let fatal = match command { + TpWorkerCommand::Ping { resp } => { + self.respond(resp, "ping", Ok(TpWorkerReply::Ack)) + } + TpWorkerCommand::RunPrefillChunks { + chunks, + sample_seed, + resp, + } => { + let result = self.execute_prefill_chunks(&chunks, sample_seed); + self.respond(resp, "prefill", result) + } + TpWorkerCommand::RunDecodeStep { + requests, + sample_seed, + resp, + } => { + let result = self.execute_decode(&requests, sample_seed); + self.respond(resp, "decode", result) + } + TpWorkerCommand::RunUnifiedStep { resp } => { + let rank = self.rank; + self.respond( + resp, + "unified step", + Err(anyhow::anyhow!( + "Qwen3.5 TP worker rank {rank} has no TP unified implementation yet" + )), + ) + } + TpWorkerCommand::DropRequest { request_id, resp } => { + self.drop_request(request_id); + self.respond(resp, "drop request", Ok(TpWorkerReply::Ack)) + } + TpWorkerCommand::Shutdown => break, + }; + if fatal { + break; + } + } + } + + fn respond( + &self, + resp: mpsc::Sender, + operation: &'static str, + result: Result, + ) -> bool { + match result { + Ok(reply) => { + let _ = resp.send(TpWorkerResponse { + rank: self.rank, + result: Ok(reply), + }); + false + } + Err(err) => { + let reason = self.poison.poison(format!( + "rank {} failed during {operation}: {err:#}", + self.rank + )); + let _ = resp.send(TpWorkerResponse { + rank: self.rank, + result: Err(anyhow::anyhow!(reason)), + }); + true + } + } + } + + fn execute_prefill_chunks( + &mut self, + chunks: &[TpPrefillChunkItem], + sample_seed: u64, + ) -> Result { + anyhow::ensure!( + !chunks.is_empty(), + "Qwen3.5 TP prefill chunk command requires at least one chunk" + ); + validate_prefill_chunks(chunks)?; + let new_requests = chunks + .iter() + .filter(|chunk| self.request_index(chunk.request_id).is_none()) + .count(); + anyhow::ensure!( + self.requests.len() + new_requests <= self.max_batch, + "Qwen3.5 TP prefill chunks would exceed worker capacity {}", + self.max_batch + ); + + let mut primary_results = Vec::new(); + let mut final_row_idx = 0usize; + for chunk in chunks { + let state_idx = self.ensure_prefill_state(chunk.request_id)?; + let state = &mut self.requests[state_idx]; + anyhow::ensure!( + state.phase == TpRequestPhase::Prefilling, + "Qwen3.5 TP request {} is already in decode state", + chunk.request_id.get() + ); + + let prompt = [chunk.prompt_tokens.as_slice()]; + let mut recurrent_refs = vec![&mut state.recurrent]; + let logits = self.model.batch_prefill_logits( + &prompt, + std::slice::from_mut(&mut state.kv), + &mut recurrent_refs, + )?; + + if chunk.finish_prefill { + if self.rank == 0 { + // TP prefill samples final chunks one row at a time. Offset + // by the final-row index so rows from the same command do + // not reuse the same sampling stream. + let row_seed = sample_seed.wrapping_add(final_row_idx as u64); + let result = self.sample_final_prefill_chunk(chunk, &logits, row_seed)?; + primary_results.push(result); + } + final_row_idx += 1; + self.requests[state_idx].phase = TpRequestPhase::Decoding; + } + } + + if self.rank == 0 { + Ok(TpWorkerReply::Prefill(PrefillResult { + requests: primary_results, + })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn sample_final_prefill_chunk( + &mut self, + chunk: &TpPrefillChunkItem, + logits: &openinfer_core::tensor::HiddenStates, + sample_seed: u64, + ) -> Result { + let cpu_logits = + snapshot_requested_logprobs(self.model.device_ctx(), logits, &[chunk.logprobs])?; + let params_refs = [&chunk.sampling_params]; + let tokens = openinfer_sample::select_batch( + self.model.device_ctx(), + logits, + ¶ms_refs, + &[0], + sample_seed, + &mut self.sample_scratch, + )?; + let first_token = tokens[0]; + let first_token_logprob = cpu_logits[0].as_ref().and_then(|row| { + openinfer_sample::token_logprob_from_row(row, first_token, chunk.logprobs) + }); + Ok(PrefillRequestResult { + request_id: chunk.request_id, + first_token, + first_token_logprob, + }) + } + + fn execute_decode( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result { + anyhow::ensure!( + !requests.is_empty(), + "Qwen3.5 TP decode command requires at least one request" + ); + validate_decode_requests(requests)?; + anyhow::ensure!( + requests.len() <= self.max_batch, + "Qwen3.5 TP decode batch {} exceeds worker capacity {}", + requests.len(), + self.max_batch + ); + + let mut primary_results = + Vec::with_capacity(if self.rank == 0 { requests.len() } else { 0 }); + for (row_idx, request) in requests.iter().enumerate() { + let state_idx = self.request_index(request.request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode request {} has no worker state", + request.request_id.get() + ) + })?; + anyhow::ensure!( + self.requests[state_idx].phase == TpRequestPhase::Decoding, + "Qwen3.5 TP request {} is not ready for decode", + request.request_id.get() + ); + + { + let state = &mut self.requests[state_idx]; + let mut kv_refs = vec![&mut state.kv]; + let mut recurrent_refs = vec![&mut state.recurrent]; + self.model.batch_decode_eager_logits( + &[request.token_id], + &mut kv_refs, + &mut recurrent_refs, + &mut self.decode_buffers, + )?; + } + + if self.rank == 0 { + let cpu_logits = snapshot_requested_logprobs( + self.model.device_ctx(), + &self.decode_buffers.logits, + &[request.logprobs], + )?; + let params_refs = [&request.sampling_params]; + let tokens = openinfer_sample::select_batch( + self.model.device_ctx(), + &self.decode_buffers.logits, + ¶ms_refs, + &[0], + sample_seed.wrapping_add(row_idx as u64), + &mut self.sample_scratch, + )?; + let token = tokens[0]; + let logprob = cpu_logits[0].as_ref().and_then(|row| { + openinfer_sample::token_logprob_from_row(row, token, request.logprobs) + }); + primary_results.push(DecodeRequestResult { + request_id: request.request_id, + token, + logprob, + }); + } + } + + if self.rank == 0 { + Ok(TpWorkerReply::Decode(DecodeResult { + requests: primary_results, + })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn ensure_prefill_state(&mut self, request_id: RequestId) -> Result { + if let Some(idx) = self.request_index(request_id) { + return Ok(idx); + } + let state = TpRequestState { + request_id, + phase: TpRequestPhase::Prefilling, + kv: self.model.alloc_kv(), + recurrent: RecurrentState::new(self.model.device_ctx(), self.model.config())?, + }; + self.requests.push(state); + Ok(self.requests.len() - 1) + } + + fn request_index(&self, request_id: RequestId) -> Option { + self.requests + .iter() + .position(|state| state.request_id == request_id) + } + + fn drop_request(&mut self, request_id: RequestId) { + if let Some(idx) = self.request_index(request_id) { + self.requests.swap_remove(idx); + } + } +} + +fn validate_prefill_chunks(chunks: &[TpPrefillChunkItem]) -> Result<()> { + let mut seen = HashSet::with_capacity(chunks.len()); + for chunk in chunks { + anyhow::ensure!( + !chunk.prompt_tokens.is_empty(), + "Qwen3.5 TP prefill chunk for request {} is empty", + chunk.request_id.get() + ); + anyhow::ensure!( + seen.insert(chunk.request_id), + "duplicate Qwen3.5 TP request id {} in one prefill chunk command", + chunk.request_id.get() + ); + } + Ok(()) +} + +fn validate_decode_requests(requests: &[TpDecodeStepItem]) -> Result<()> { + let mut seen = HashSet::with_capacity(requests.len()); + for request in requests { + anyhow::ensure!( + seen.insert(request.request_id), + "duplicate Qwen3.5 TP request id {} in one decode command", + request.request_id.get() + ); + } + Ok(()) +} + +impl From for TpPrefillChunkItem { + fn from(request: PrefillStepItem) -> Self { + Self::new( + request.request_id, + request.prompt_tokens, + request.logprobs, + true, + ) + } +} + +impl From for TpDecodeStepItem { + fn from(request: DecodeStepItem) -> Self { + Self::new( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + ) + } +} + +fn wait_for_acks( + responses: mpsc::Receiver, + expected: usize, + op_name: &'static str, + poison: &TpRuntimePoison, +) -> Result<()> { + for _ in 0..expected { + let response = recv_runtime_response(&responses, op_name, poison)?; + match response.result? { + TpWorkerReply::Ack => {} + TpWorkerReply::Prefill(_) => { + anyhow::bail!("Qwen3.5 TP {op_name} unexpectedly returned prefill result") + } + TpWorkerReply::Decode(_) => { + anyhow::bail!("Qwen3.5 TP {op_name} unexpectedly returned decode result") + } + } + } + Ok(()) +} + +fn wait_for_prefill( + responses: mpsc::Receiver, + expected: usize, + poison: &TpRuntimePoison, +) -> Result { + let mut result = None; + for _ in 0..expected { + let response = recv_runtime_response(&responses, "prefill", poison)?; + match response.result? { + TpWorkerReply::Ack => {} + TpWorkerReply::Prefill(prefill) => { + anyhow::ensure!( + response.rank == 0, + "Qwen3.5 TP prefill returned a primary result from rank {}", + response.rank + ); + anyhow::ensure!( + result.is_none(), + "Qwen3.5 TP prefill returned multiple primary results" + ); + result = Some(prefill); + } + TpWorkerReply::Decode(_) => { + anyhow::bail!("Qwen3.5 TP prefill unexpectedly returned decode result") + } + } + } + result.ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP prefill returned no primary result")) +} + +fn wait_for_decode( + responses: mpsc::Receiver, + expected: usize, + poison: &TpRuntimePoison, +) -> Result { + let mut result = None; + for _ in 0..expected { + let response = recv_runtime_response(&responses, "decode", poison)?; + match response.result? { + TpWorkerReply::Ack => {} + TpWorkerReply::Decode(decode) => { + anyhow::ensure!( + response.rank == 0, + "Qwen3.5 TP decode returned a primary result from rank {}", + response.rank + ); + anyhow::ensure!( + result.is_none(), + "Qwen3.5 TP decode returned multiple primary results" + ); + result = Some(decode); + } + TpWorkerReply::Prefill(_) => { + anyhow::bail!("Qwen3.5 TP decode unexpectedly returned prefill result") + } + } + } + result.ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP decode returned no primary result")) +} + +fn recv_runtime_response( + responses: &mpsc::Receiver, + operation: &'static str, + poison: &TpRuntimePoison, +) -> Result { + match responses.recv_timeout(TP_RUNTIME_STEP_TIMEOUT) { + Ok(response) => Ok(response), + Err(mpsc::RecvTimeoutError::Disconnected) => { + let reason = poison.poison(format!("response channel disconnected during {operation}")); + Err(anyhow::anyhow!(reason)) + } + Err(mpsc::RecvTimeoutError::Timeout) => fatal_tp_abort(&format!( + "Qwen3.5 TP {operation} did not complete within {}s", + TP_RUNTIME_STEP_TIMEOUT.as_secs() + )), + } +} + +fn fatal_tp_abort(message: &str) -> ! { + eprintln!("{message}; aborting"); + log::error!("{message}; aborting"); + std::process::abort(); +} + +struct CublasThreadGuard; + +impl Drop for CublasThreadGuard { + fn drop(&mut self) { + unsafe { + crate::ffi::cublas_destroy(); + } + } +} + +fn bind_worker_thread(model: &Qwen35Model) -> Result { + let ctx = model.device_ctx(); + unsafe { + let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); + if err != 0 { + return Err(anyhow::anyhow!( + "Failed to set CUDA device {} on Qwen3.5 TP worker thread: cudaError={}", + ctx.device_ordinal, + err + )); + } + } + ctx.ctx.bind_to_thread().map_err(|e| { + anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 TP worker thread: {e}") + })?; + unsafe { + crate::ffi::cublas_init(); + } + Ok(CublasThreadGuard) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn startup_gate_cancel_releases_waiting_workers() { + let gate = Arc::new(TpStartupGate::default()); + let worker_gate = Arc::clone(&gate); + let (done_tx, done_rx) = mpsc::channel(); + let waiter = thread::spawn(move || { + let _ = done_tx.send(worker_gate.wait()); + }); + + gate.cancel(); + + assert_eq!( + done_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("cancelled startup gate should release workers within one second"), + false + ); + waiter.join().unwrap(); + } + + #[test] + fn nccl_startup_watchdog_disarms_after_success() { + let (done_tx, watchdog) = spawn_nccl_startup_watchdog().unwrap(); + disarm_nccl_startup_watchdog(done_tx, watchdog).unwrap(); + } + + #[test] + fn runtime_poison_preserves_first_failure() { + let poison = TpRuntimePoison::default(); + assert_eq!(poison.poison("rank 1 OOM".into()), "rank 1 OOM"); + assert_eq!(poison.poison("rank 0 NCCL error".into()), "rank 1 OOM"); + let err = poison.ensure_healthy().unwrap_err().to_string(); + assert!(err.contains("rank 1 OOM")); + assert!(!err.contains("rank 0 NCCL error")); + } + + #[test] + fn runtime_response_reports_any_rank_failure_immediately() { + let poison = TpRuntimePoison::default(); + let (tx, rx) = mpsc::channel(); + tx.send(TpWorkerResponse { + rank: 1, + result: Err(anyhow::anyhow!("rank 1 failed")), + }) + .unwrap(); + + let err = wait_for_acks(rx, 2, "test", &poison) + .unwrap_err() + .to_string(); + assert!(err.contains("rank 1 failed")); + } + + #[test] + fn disconnected_runtime_response_poisons_executor() { + let poison = TpRuntimePoison::default(); + let (tx, rx) = mpsc::channel(); + drop(tx); + + let err = recv_runtime_response(&rx, "test", &poison) + .unwrap_err() + .to_string(); + assert!(err.contains("response channel disconnected during test")); + assert!(poison.ensure_healthy().is_err()); + } + + #[test] + fn prefill_scratch_tokens_follow_budget_and_chunk_cap() { + assert_eq!(prefill_scratch_tokens(1_024), 1_024); + assert_eq!(prefill_scratch_tokens(PREFILL_CHUNK_LEN), 20_000); + assert_eq!(prefill_scratch_tokens(40_000), 20_000); + } + + #[test] + fn recurrent_capacity_reserves_runtime_and_prefill_headroom() { + const MIB: usize = 1024 * 1024; + assert_eq!( + effective_recurrent_capacity(64, 10_000 * MIB, 50 * MIB, 512 * MIB, 1_000 * MIB,), + 64 + ); + assert_eq!( + effective_recurrent_capacity(64, 2_061 * MIB, 50 * MIB, 512 * MIB, 1_000 * MIB,), + 10 + ); + assert_eq!( + effective_recurrent_capacity(64, 1_511 * MIB, 50 * MIB, 512 * MIB, 1_000 * MIB,), + 0 + ); + } + + #[test] + fn zero_sized_recurrent_state_keeps_requested_capacity() { + assert_eq!( + effective_recurrent_capacity(64, 0, 0, usize::MAX, usize::MAX), + 64 + ); + } + + #[test] + fn limits_constructor_rejects_zero_prefill_budget_before_loading() { + let err = match Qwen35TpExecutor::from_runtime_with_limits("unused", false, &[0, 1], 1, 0) { + Ok(_) => panic!("zero TP prefill budget should fail"), + Err(err) => err.to_string(), + }; + assert!(err.contains("max_prefill_tokens must be positive")); + } + + #[test] + fn rejects_single_device_topology() { + let err = match Qwen35TpExecutor::from_runtime_with_capacity("unused", false, &[0], 1) { + Ok(_) => panic!("single-device TP topology should fail"), + Err(err) => err.to_string(), + }; + assert!(err.contains("requires at least two CUDA devices")); + } + + #[test] + fn rejects_tensor_parallel_cuda_graph() { + let err = match Qwen35TpExecutor::from_runtime_with_capacity("unused", true, &[0, 1], 1) { + Ok(_) => panic!("TP CUDA Graph should fail"), + Err(err) => err.to_string(), + }; + assert!(err.contains("eager execution only")); + } + + #[test] + fn validates_prefill_chunk_shape() { + let empty = [TpPrefillChunkItem::new(RequestId::new(1), vec![], 0, false)]; + let err = validate_prefill_chunks(&empty).unwrap_err().to_string(); + assert!(err.contains("is empty")); + + let duplicate = [ + TpPrefillChunkItem::new(RequestId::new(1), vec![151_646], 0, false), + TpPrefillChunkItem::new(RequestId::new(1), vec![9707], 0, true), + ]; + let err = validate_prefill_chunks(&duplicate).unwrap_err().to_string(); + assert!(err.contains("duplicate")); + } + + #[test] + fn validates_decode_request_shape() { + validate_decode_requests(&[TpDecodeStepItem::new( + RequestId::new(1), + 9707, + 0, + SamplingParams::default(), + )]) + .expect("single decode request is valid"); + + let duplicate = [ + TpDecodeStepItem::new(RequestId::new(1), 9707, 0, SamplingParams::default()), + TpDecodeStepItem::new(RequestId::new(1), 560, 0, SamplingParams::default()), + ]; + let err = validate_decode_requests(&duplicate) + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate")); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn starts_tp2_workers_and_broadcasts_lifecycle_commands() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + assert_eq!(executor.world_size(), 2); + assert_eq!(executor.max_batch(), 1); + executor.ping_all().expect("ping all workers"); + executor + .drop_request(RequestId::new(7)) + .expect("drop request"); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_default_capacity_is_memory_safe() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let executor = Qwen35TpExecutor::from_runtime(&model_path, false, &[0, 1]) + .expect("start TP2 executor with memory-derived capacity"); + eprintln!( + "Qwen3.5 TP2 memory-derived max_batch={}", + executor.max_batch() + ); + assert!(executor.max_batch() > 0); + assert!(executor.max_batch() <= MAX_BATCH); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_prefill_runs_and_returns_primary_result() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + let request_id = RequestId::new(11); + let request = PrefillStepItem::new(request_id, vec![151_646, 9707], 0); + let result = executor + .execute_prefill(PrefillPlan { + requests: &[request], + }) + .expect("run TP2 prefill"); + assert_eq!(result.requests.len(), 1); + assert_eq!(result.requests[0].request_id, request_id); + executor + .drop_request(request_id) + .expect("drop prefetched request"); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_chunked_prefill_advances_existing_request_state() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + let request_id = RequestId::new(13); + let first = TpPrefillChunkItem::new(request_id, vec![151_646], 0, false); + let first_result = executor + .execute_prefill_chunks(&[first]) + .expect("run non-final TP2 prefill chunk"); + assert!(first_result.requests.is_empty()); + + let final_chunk = TpPrefillChunkItem::new(request_id, vec![9707], 0, true); + let final_result = executor + .execute_prefill_chunks(&[final_chunk]) + .expect("run final TP2 prefill chunk"); + assert_eq!(final_result.requests.len(), 1); + assert_eq!(final_result.requests[0].request_id, request_id); + + executor + .drop_request(request_id) + .expect("drop chunk-prefilled request"); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_decode_runs_after_prefill() { + let model_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + let request_id = RequestId::new(17); + let request = PrefillStepItem::new(request_id, vec![151_646, 9707], 0); + let prefill = executor + .execute_prefill(PrefillPlan { + requests: &[request], + }) + .expect("run TP2 prefill"); + assert_eq!(prefill.requests.len(), 1); + assert_eq!(prefill.requests[0].request_id, request_id); + + let decode_request = DecodeStepItem::new(request_id, prefill.requests[0].first_token, 0); + let decode = executor + .execute_decode(DecodePlan { + requests: &[decode_request], + }) + .expect("run TP2 eager decode"); + assert_eq!(decode.requests.len(), 1); + assert_eq!(decode.requests[0].request_id, request_id); + + executor + .drop_request(request_id) + .expect("drop decoded request"); + } +} diff --git a/openinfer-qwen35-4b/src/weights.rs b/openinfer-qwen35-4b/src/weights.rs index 617116b5d..b5dd13d59 100644 --- a/openinfer-qwen35-4b/src/weights.rs +++ b/openinfer-qwen35-4b/src/weights.rs @@ -1,13 +1,16 @@ use anyhow::Result; use cudarc::driver::CudaSlice; +use cudarc::nccl::safe::{Comm, ReduceOp}; use log::{debug, info}; +use safetensors::SafeTensors; +use std::collections::HashMap; use std::time::Instant; -use super::config::{Config35, LayerType}; -use openinfer_core::tensor::{DeviceContext, DeviceMatrix, DeviceVec}; +use super::config::{Config35, LayerType, TensorParallelConfig}; +use openinfer_core::tensor::{DeviceContext, DeviceMatrix, DeviceVec, HiddenStates}; use openinfer_core::weight_loader::{ deserialize_shards, load_shard_info_fixed, load_tensor_1d, load_tensor_1d_f32, load_tensor_2d, - mmap_shards, precompute_rope, + load_tensor_2d_col_shard, load_tensor_2d_row_shard, mmap_shards, precompute_rope, }; /// Full attention layer weights (8 layers in Qwen3.5-4B). @@ -68,10 +71,28 @@ pub(super) struct TransformerBlock35 { pub(super) mlp: MLP35, } +#[derive(Clone, Copy, Debug)] +pub(crate) struct ModelRuntimeConfig { + pub(crate) enable_cuda_graph: bool, + pub(crate) tensor_parallel: Option, + pub(crate) device_ordinal: usize, +} + +impl Default for ModelRuntimeConfig { + fn default() -> Self { + Self { + enable_cuda_graph: true, + tensor_parallel: None, + device_ordinal: 0, + } + } +} + /// Qwen3.5 model (text-only). pub struct Qwen35Model { pub(super) ctx: DeviceContext, pub(super) config: Config35, + pub(super) tensor_parallel: TensorParallelConfig, pub(super) embed_tokens: DeviceMatrix, pub(super) lm_head: Option, pub(super) layers: Vec, @@ -83,13 +104,51 @@ pub struct Qwen35Model { pub(super) kv_pool: openinfer_core::kv_pool::KvPool, /// Decode-slot count the recurrent-state reserve was sized for. pub(super) reserved_decode_slots: usize, + pub(super) tp_comm: Option, } +// SAFETY: A Qwen3.5 model instance is bound to one CUDA device and driven from +// one owning scheduler/worker thread at a time. TP constructs one independent +// rank-local model per worker; the model is moved between threads only during +// startup, never shared for concurrent mutation. +unsafe impl Send for Qwen35Model {} +unsafe impl Sync for Qwen35Model {} + /// Graph slot state + one in-flight prefill transient per decode slot. const STATES_PER_DECODE_SLOT: usize = 2; /// KV-pool floor, also the low-memory fail-fast threshold. const MIN_KV_PAGES: usize = 64; +impl Qwen35Model { + pub fn from_safetensors_with_options( + model_path: &str, + enable_cuda_graph: bool, + ) -> Result { + Self::from_safetensors_with_runtime( + model_path, + ModelRuntimeConfig { + enable_cuda_graph, + ..Default::default() + }, + ) + } + + pub fn from_safetensors_with_device_options( + model_path: &str, + enable_cuda_graph: bool, + device_ordinal: usize, + ) -> Result { + Self::from_safetensors_with_runtime( + model_path, + ModelRuntimeConfig { + enable_cuda_graph, + device_ordinal, + ..Default::default() + }, + ) + } +} + impl Qwen35Model { /// `max_batch` must be a decode bucket ({1,2,4,8,16,32,64}). pub fn from_safetensors( @@ -102,18 +161,48 @@ impl Qwen35Model { "decode batch capacity must be one of {:?}, got {max_batch}", super::batch_decode_graph::BATCH_BUCKETS, ); + Self::from_safetensors_with_runtime_and_capacity( + model_path, + ModelRuntimeConfig { + device_ordinal, + ..Default::default() + }, + max_batch, + ) + } + + pub(crate) fn from_safetensors_with_runtime( + model_path: &str, + runtime: ModelRuntimeConfig, + ) -> Result { + Self::from_safetensors_with_runtime_and_capacity( + model_path, + runtime, + super::batch_decode_graph::MAX_BATCH, + ) + } + + fn from_safetensors_with_runtime_and_capacity( + model_path: &str, + runtime: ModelRuntimeConfig, + max_batch: usize, + ) -> Result { info!("Loading Qwen3.5 model from: {}", model_path); - debug!("Initializing GPU"); - let ctx = DeviceContext::new_with_device(device_ordinal)?; + debug!("Initializing GPU device {}", runtime.device_ordinal); + let ctx = DeviceContext::new_with_device(runtime.device_ordinal)?; let mut config = Config35::from_file(model_path)?; + let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); + tensor_parallel.validate_for(&config, runtime.enable_cuda_graph)?; debug!( - "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}", + "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}, tp_rank={}, tp_world_size={}", config.hidden_size, config.num_hidden_layers, config.num_full_attention_layers(), config.num_hidden_layers - config.num_full_attention_layers(), - config.max_position_embeddings + config.max_position_embeddings, + tensor_parallel.rank, + tensor_parallel.world_size, ); let effective_vocab = super::config::tokenizer_effective_vocab(model_path)?; anyhow::ensure!( @@ -173,6 +262,9 @@ impl Qwen35Model { config.num_hidden_layers ); let mut layers = Vec::with_capacity(config.num_hidden_layers); + let (_, q_rows) = tensor_parallel.shard_range(config.full_attn_q_dim()); + let (kv_row_offset, kv_rows) = tensor_parallel.shard_range(config.full_attn_kv_dim()); + let (inter_row_offset, inter_rows) = tensor_parallel.shard_range(config.intermediate_size); for i in 0..config.num_hidden_layers { let prefix = format!("{}.layers.{}", wp, i); let layer_type = config.layer_types[i]; @@ -181,29 +273,40 @@ impl Qwen35Model { LayerType::FullAttention => { let attn_prefix = format!("{}.self_attn", prefix); LayerKind::FullAttention(FullAttentionLayer { - q_proj: load_tensor_2d( + q_proj: load_full_attention_gated_q_proj( &ctx, &shards, &weight_map, &format!("{}.q_proj.weight", attn_prefix), + &config, + tensor_parallel, )?, - k_proj: load_tensor_2d( + k_proj: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.k_proj.weight", attn_prefix), + tensor_parallel, + kv_row_offset, + kv_rows, )?, - v_proj: load_tensor_2d( + v_proj: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.v_proj.weight", attn_prefix), + tensor_parallel, + kv_row_offset, + kv_rows, )?, - o_proj: load_tensor_2d( + o_proj: load_tensor_2d_col_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.o_proj.weight", attn_prefix), + tensor_parallel, + tensor_parallel.shard_range(config.full_attn_q_dim()).0, + q_rows, )?, q_norm: load_tensor_1d( &ctx, @@ -280,17 +383,23 @@ impl Qwen35Model { } }; - let gate_proj = load_tensor_2d( + let gate_proj = load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.mlp.gate_proj.weight", prefix), + tensor_parallel, + inter_row_offset, + inter_rows, )?; - let up_proj = load_tensor_2d( + let up_proj = load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.mlp.up_proj.weight", prefix), + tensor_parallel, + inter_row_offset, + inter_rows, )?; let gate_up_proj = DeviceMatrix::vstack(&ctx, &[&gate_proj, &up_proj])?; drop(gate_proj); @@ -312,11 +421,14 @@ impl Qwen35Model { )?, mlp: MLP35 { gate_up_proj, - down_proj: load_tensor_2d( + down_proj: load_tensor_2d_col_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.mlp.down_proj.weight", prefix), + tensor_parallel, + inter_row_offset, + inter_rows, )?, }, }; @@ -348,12 +460,17 @@ impl Qwen35Model { "GPU model loaded in {:.0}ms", t_gpu.elapsed().as_secs_f64() * 1e3 ); + if runtime.enable_cuda_graph { + debug!("Decode path CUDA Graph is enabled"); + } else { + debug!("Decode path CUDA Graph is disabled"); + } // Paged KV pool for the 8 full-attention layers. let page_size = 16usize; let num_full_layers = config.num_full_attention_layers(); let layout = openinfer_core::kv_pool::KvLayout::new( num_full_layers, - config.num_key_value_heads, + config.local_num_key_value_heads(tensor_parallel), config.head_dim, page_size, ); @@ -393,7 +510,7 @@ impl Qwen35Model { let kv_pool = openinfer_core::kv_pool::KvPool::new( &ctx, num_full_layers, - config.num_key_value_heads, + config.local_num_key_value_heads(tensor_parallel), config.head_dim, page_size, num_pages, @@ -402,6 +519,7 @@ impl Qwen35Model { Ok(Self { ctx, config, + tensor_parallel, embed_tokens, lm_head, layers, @@ -410,6 +528,7 @@ impl Qwen35Model { sin_cache, kv_pool, reserved_decode_slots: max_batch, + tp_comm: None, }) } @@ -443,6 +562,22 @@ impl Qwen35Model { &self.kv_pool } + pub(crate) fn attach_tp_comm(&mut self, comm: Comm) { + self.tp_comm = Some(comm); + } + + pub(crate) fn all_reduce_hidden(&self, hidden: &mut HiddenStates) -> Result<()> { + self.all_reduce_hidden_untraced(hidden) + } + + pub(crate) fn all_reduce_hidden_untraced(&self, hidden: &mut HiddenStates) -> Result<()> { + if let Some(comm) = &self.tp_comm { + comm.all_reduce_in_place(&mut hidden.data, &ReduceOp::Sum) + .map_err(|e| anyhow::anyhow!("Qwen3.5 NCCL all-reduce failed: {e:?}"))?; + } + Ok(()) + } + /// Tune small-batch decode GEMM algorithms on the thread that will capture /// or replay the CUDA Graph. cuBLASLt plans are thread-local, so scheduler /// workers and model-local executors must call this after binding CUDA. @@ -452,12 +587,13 @@ impl Qwen35Model { let ctx = &self.ctx; let hidden = self.config.hidden_size; let vocab = self.config.selection_vocab; - let full_q = self.config.full_attn_q_proj_dim(); - let full_kv = self.config.full_attn_kv_dim(); + let tp = self.tensor_parallel; + let full_q = self.config.local_full_attn_gated_q_dim(tp); + let full_kv = self.config.local_full_attn_kv_dim(tp); let linear_qkv = self.config.linear_attn_qkv_dim(); let linear_z = self.config.linear_attn_z_dim(); let linear_ba = self.config.linear_num_value_heads; - let intermediate = self.config.intermediate_size; + let intermediate = self.config.local_intermediate_size(tp); let full_q_samples: Vec<_> = self .layers @@ -555,11 +691,38 @@ impl Qwen35Model { pub(crate) fn create_batch_decode_graph_state( &self, ) -> anyhow::Result { + self.create_batch_decode_graph_state_with_capacity(self.reserved_decode_slots) + } + + pub(crate) fn create_batch_decode_graph_state_with_capacity( + &self, + max_batch: usize, + ) -> anyhow::Result { + anyhow::ensure!( + max_batch <= self.reserved_decode_slots, + "requested graph capacity {max_batch} exceeds loaded capacity {}", + self.reserved_decode_slots + ); super::batch_decode_graph::BatchDecodeGraphState::with_capacity( &self.ctx, &self.config, + self.tensor_parallel, &self.kv_pool, - self.reserved_decode_slots, + max_batch, + ) + } + + pub(crate) fn create_batch_decode_buffers_with_capacity( + &self, + max_batch: usize, + ) -> anyhow::Result { + super::decode_buffers::BatchDecodeBuffers35::new( + &self.ctx, + &self.config, + self.tensor_parallel, + max_batch, + self.kv_pool.capacity_pages(), + self.kv_pool.padding_page_id(), ) } @@ -579,3 +742,154 @@ fn tune_if_nonempty( } crate::ops::gemm_lt_tune(ctx, samples, rows, n) } + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct GatedQShardRange { + row_offset: usize, + rows: usize, +} + +fn full_attention_gated_q_shard_range( + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> GatedQShardRange { + // HF/OpenInfer kernels interpret q_proj rows as per-head [q, gate] chunks. + // Keep each local head's q rows adjacent to its gate rows. + let local_heads = config.local_num_attention_heads(tensor_parallel); + let head_start = tensor_parallel.rank * local_heads; + GatedQShardRange { + row_offset: head_start * config.head_dim * 2, + rows: local_heads * config.head_dim * 2, + } +} + +fn load_full_attention_gated_q_proj( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> Result { + if !tensor_parallel.is_sharded() { + return load_tensor_2d(ctx, shards, weight_map, name); + } + + let range = full_attention_gated_q_shard_range(config, tensor_parallel); + load_tensor_2d_row_shard(ctx, shards, weight_map, name, range.row_offset, range.rows) +} + +fn load_tensor_2d_row_shard_if_needed( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + tensor_parallel: TensorParallelConfig, + row_offset: usize, + rows: usize, +) -> Result { + if tensor_parallel.is_sharded() { + load_tensor_2d_row_shard(ctx, shards, weight_map, name, row_offset, rows) + } else { + load_tensor_2d(ctx, shards, weight_map, name) + } +} + +fn load_tensor_2d_col_shard_if_needed( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + tensor_parallel: TensorParallelConfig, + col_offset: usize, + cols: usize, +) -> Result { + if tensor_parallel.is_sharded() { + load_tensor_2d_col_shard(ctx, shards, weight_map, name, col_offset, cols) + } else { + load_tensor_2d(ctx, shards, weight_map, name) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_config() -> Config35 { + Config35 { + hidden_size: 2560, + intermediate_size: 9216, + num_hidden_layers: 32, + vocab_size: 248320, + selection_vocab: 248320, + rms_norm_eps: 1e-6, + eos_token_id: 151645, + num_attention_heads: 16, + num_key_value_heads: 4, + head_dim: 256, + linear_num_key_heads: 16, + linear_key_head_dim: 128, + linear_num_value_heads: 32, + linear_value_head_dim: 128, + linear_conv_kernel_dim: 4, + rope_theta: 10_000.0, + rotary_dim: 64, + max_position_embeddings: 262_144, + tie_word_embeddings: true, + layer_types: vec![LayerType::LinearAttention; 32], + } + } + + #[test] + fn gated_q_shard_range_keeps_matching_q_and_gate_rows() { + let config = test_config(); + + let rank0 = full_attention_gated_q_shard_range( + &config, + TensorParallelConfig { + rank: 0, + world_size: 2, + }, + ); + assert_eq!( + rank0, + GatedQShardRange { + row_offset: 0, + rows: 4096, + } + ); + + let rank1 = full_attention_gated_q_shard_range( + &config, + TensorParallelConfig { + rank: 1, + world_size: 2, + }, + ); + assert_eq!( + rank1, + GatedQShardRange { + row_offset: 4096, + rows: 4096, + } + ); + } + + #[test] + fn mlp_tp2_uses_matching_gate_up_rows_and_down_cols() { + let config = test_config(); + let tp = TensorParallelConfig { + rank: 1, + world_size: 2, + }; + + let (inter_offset, inter_rows) = tp.shard_range(config.intermediate_size); + assert_eq!((inter_offset, inter_rows), (4608, 4608)); + assert_eq!(config.local_intermediate_size(tp), inter_rows); + + let local_gate_up_rows = 2 * inter_rows; + let local_down_cols = inter_rows; + assert_eq!(local_gate_up_rows, 9216); + assert_eq!(local_down_cols, 4608); + } +} diff --git a/openinfer-qwen35-4b/tests/common/mod.rs b/openinfer-qwen35-4b/tests/common/mod.rs index 6614d672d..2b14de6dd 100644 --- a/openinfer-qwen35-4b/tests/common/mod.rs +++ b/openinfer-qwen35-4b/tests/common/mod.rs @@ -1,6 +1,36 @@ use vllm_text::tokenizer::DynTokenizer; +#[allow(dead_code)] pub(crate) fn load_tokenizer(model_path: &str) -> DynTokenizer { openinfer_vllm_support::load_tokenizer(model_path) .unwrap_or_else(|err| panic!("Failed to load tokenizer for {model_path}: {err}")) } + +pub(crate) fn tp2_device_ordinals() -> Vec { + const ENV: &str = "OPENINFER_TEST_TP_DEVICES"; + let value = match std::env::var(ENV) { + Ok(value) => value, + Err(_) => return vec![0, 1], + }; + + let devices: Vec = value + .split(',') + .map(str::trim) + .filter(|part| !part.is_empty()) + .map(|part| { + part.parse::() + .unwrap_or_else(|err| panic!("{ENV} must be comma-separated CUDA ordinals: {err}")) + }) + .collect(); + + assert_eq!( + devices.len(), + 2, + "{ENV} must specify exactly two CUDA ordinals for TP2, e.g. 0,1 or 2,3" + ); + assert_ne!( + devices[0], devices[1], + "{ENV} must specify two distinct CUDA ordinals for TP2" + ); + devices +} diff --git a/openinfer-qwen35-4b/tests/e2e_scheduler.rs b/openinfer-qwen35-4b/tests/e2e_scheduler.rs index c7fba6e38..655d684e7 100644 --- a/openinfer-qwen35-4b/tests/e2e_scheduler.rs +++ b/openinfer-qwen35-4b/tests/e2e_scheduler.rs @@ -2,15 +2,13 @@ /// /// Tests the Qwen3.5 reduced-capacity scheduler path (batch prefill + /// CUDA Graph decode) with sequential, concurrent, and consumer-drop requests. -use std::collections::HashSet; -use std::time::Instant; +use std::{collections::HashSet, path::Path, time::Instant}; use log::info; -use openinfer_core::engine::FinishReason; use openinfer_core::engine::{ - EngineHandle, EngineLoadOptions, GenerateRequest, TokenEvent, TokenLogprob, TokenSink, - TokenStreamReceiver, + EngineHandle, EngineLoadOptions, FinishReason, GenerateRequest, TokenEvent, TokenLogprob, + TokenSink, TokenStreamReceiver, }; use openinfer_core::sampler::SamplingParams; use vllm_text::tokenizer::DynTokenizer; @@ -321,33 +319,17 @@ fn assert_no_model_wide_collapse(collapses: &[(&str, Collapse)]) { } } -#[test] -fn test_e2e_qwen35_scheduler() { +fn run_full_scheduler_e2e( + handle: &EngineHandle, + tokenizer: &DynTokenizer, + max_context_tokens: usize, + label: &str, +) { // logging intentionally left to the test harness - let model_path = get_model_path(); - let max_context_tokens = max_position_embeddings(&model_path); - - info!("Loading Qwen3.5 model for scheduler test..."); - let start = Instant::now(); - let tokenizer = common::load_tokenizer(&model_path); - let handle = openinfer_qwen35_4b::start_engine( - std::path::Path::new(&model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: vec![0], - seed: 42, - ..EngineLoadOptions::default() - }, - 8, - openinfer_qwen35_4b::DEFAULT_MAX_PREFILL_TOKENS, - ) - .expect("Failed to start Qwen3.5 scheduler"); - info!("scheduler loaded in {:.2?}", start.elapsed()); - // ── 0. Static context-window rejection ───────────────────────────── info!("=== Phase 0: Context-window rejection ==="); - expect_context_window_rejection(&handle, max_context_tokens); + expect_context_window_rejection(handle, max_context_tokens); info!(" PASS: over-context request rejected before prefill"); // ── 1. logprobs must not change greedy tokens ───────────────────── @@ -355,9 +337,9 @@ fn test_e2e_qwen35_scheduler() { for case in CASES.iter().take(3) { let max_tokens = case.max_new_tokens.min(16); let no_logprobs = - generate_tokens_with_logprobs(&handle, &tokenizer, case.prompt, max_tokens, 0); + generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 0); let with_logprobs = - generate_tokens_with_logprobs(&handle, &tokenizer, case.prompt, max_tokens, 1); + generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 1); assert_eq!(no_logprobs.finish_reason, with_logprobs.finish_reason); assert_eq!( no_logprobs.tokens, with_logprobs.tokens, @@ -392,7 +374,7 @@ fn test_e2e_qwen35_scheduler() { info!("--- {:?} ---", case.name); let start = Instant::now(); let (tokens, finish_reason) = - generate_tokens(&handle, &tokenizer, case.prompt, case.max_new_tokens); + generate_tokens(handle, tokenizer, case.prompt, case.max_new_tokens); let elapsed = start.elapsed(); let text = tokenizer.decode(&tokens, true).expect("decode failed"); @@ -419,7 +401,7 @@ fn test_e2e_qwen35_scheduler() { // ── 3. Multi-request (scheduler state reuse) ──────────────────────── info!("=== Phase 3: Multi-request ==="); for case in CASES { - let (tokens, _) = generate_tokens(&handle, &tokenizer, case.prompt, case.max_new_tokens); + let (tokens, _) = generate_tokens(handle, tokenizer, case.prompt, case.max_new_tokens); let text = tokenizer.decode(&tokens, true).expect("decode failed"); assert!( !text.is_empty(), @@ -539,10 +521,68 @@ fn test_e2e_qwen35_scheduler() { } // Verify scheduler survives - let (tokens, _) = generate_tokens(&handle, &tokenizer, "Hello", 5); + let (tokens, _) = generate_tokens(handle, tokenizer, "Hello", 5); let text = tokenizer.decode(&tokens, true).expect("decode failed"); assert!(!text.is_empty(), "scheduler dead after consumer drop"); info!(" PASS: scheduler survived consumer drop"); - info!("All Qwen3.5 scheduler tests passed!"); + info!("All Qwen3.5 scheduler tests passed for {label}!"); +} + +fn context_limit_for(handle: &EngineHandle, model_path: &str) -> usize { + handle + .servable_len() + .map(|len| len as usize) + .unwrap_or_else(|| max_position_embeddings(model_path)) +} + +#[test] +fn test_e2e_qwen35_scheduler() { + let model_path = get_model_path(); + + info!("Loading Qwen3.5 model for scheduler test..."); + let start = Instant::now(); + let model = + openinfer_qwen35_4b::runtime::Qwen35Model::from_safetensors_with_options(&model_path, true) + .expect("Failed to load model"); + let tokenizer = common::load_tokenizer(&model_path); + // Use reduced batch capacity (8) to fit on 16GB GPUs alongside the model. + let handle = openinfer_qwen35_4b::runtime::start_with_capacity( + model, + 42, + 8, + openinfer_qwen35_4b::DEFAULT_MAX_PREFILL_TOKENS, + ) + .expect("Failed to start Qwen3.5 scheduler"); + info!("scheduler loaded in {:.2?}", start.elapsed()); + + let max_context_tokens = context_limit_for(&handle, &model_path); + run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP1"); +} + +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn test_e2e_qwen35_scheduler_tp2() { + let model_path = get_model_path(); + + info!("Loading Qwen3.5 TP2 model for scheduler test..."); + let start = Instant::now(); + let tokenizer = common::load_tokenizer(&model_path); + // TP Phase 1 is eager-only; CUDA Graph must stay disabled for multi-device startup. + let handle = openinfer_qwen35_4b::start_engine_with_capacity( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: false, + device_ordinals: common::tp2_device_ordinals(), + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + openinfer_qwen35_4b::DEFAULT_MAX_PREFILL_TOKENS, + ) + .expect("Failed to start Qwen3.5 TP2 scheduler"); + info!("TP2 scheduler loaded in {:.2?}", start.elapsed()); + + let max_context_tokens = context_limit_for(&handle, &model_path); + run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2"); } diff --git a/openinfer-qwen35-4b/tests/hf_golden_gate.rs b/openinfer-qwen35-4b/tests/hf_golden_gate.rs index 7fb1712b7..14bf81113 100644 --- a/openinfer-qwen35-4b/tests/hf_golden_gate.rs +++ b/openinfer-qwen35-4b/tests/hf_golden_gate.rs @@ -17,11 +17,14 @@ use std::path::{Path, PathBuf}; use openinfer_core::engine::TokenLogprob; use openinfer_qwen35_4b::runtime::{ - DecodePlan, DecodeStepItem, PrefillPlan, PrefillStepItem, Qwen35Executor, RequestId, + DecodePlan, DecodeStepItem, PrefillPlan, PrefillStepItem, Qwen35Executor, Qwen35TpExecutor, + RequestId, }; use safetensors::{Dtype, SafeTensors}; use sha2::{Digest, Sha256}; +mod common; + const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); const GOLDEN_ENV: &str = "OPENINFER_QWEN35_HF_GOLDEN"; const LONG_GOLDEN_ENV: &str = "OPENINFER_QWEN35_HF_LONG_GOLDEN"; @@ -514,6 +517,91 @@ fn run(g: &Golden, ex: &mut Qwen35Executor, seqs: &[usize], batched: bool) -> (S (stats, fingerprint) } +fn run_tp(g: &Golden, ex: &Qwen35TpExecutor, seqs: &[usize], batched: bool) -> (Stats, Vec) { + let mut stats = Stats::default(); + let mut fingerprint = Vec::new(); + let mut fold = |stats: &mut Stats, seq, pos, pega: &[(u32, f32)]| { + fingerprint.push(pega[0].1); + check_position(stats, seq, pos, pega, &g.topk(seq, pos)); + }; + + if batched { + let ids: Vec = seqs + .iter() + .map(|&seq| RequestId::new(10_000 + seq as u64)) + .collect(); + let items: Vec = seqs + .iter() + .zip(&ids) + .map(|(&seq, &id)| prefill_item(id, g.prompt(seq))) + .collect(); + let pr = ex + .execute_prefill(PrefillPlan { requests: &items }) + .expect("TP2 prefill"); + for (i, &seq) in seqs.iter().enumerate() { + fold( + &mut stats, + seq, + 0, + &top_logprobs(pr.requests[i].first_token_logprob.as_ref()), + ); + } + + for step in 0..g.decode_len { + let items: Vec = seqs + .iter() + .zip(&ids) + .map(|(&seq, &id)| decode_item(id, g.decode(seq, step))) + .collect(); + let dr = ex + .execute_decode(DecodePlan { requests: &items }) + .expect("TP2 decode"); + for (i, &seq) in seqs.iter().enumerate() { + fold( + &mut stats, + seq, + step + 1, + &top_logprobs(dr.requests[i].logprob.as_ref()), + ); + } + } + + for &id in &ids { + ex.drop_request(id).expect("TP2 drop request"); + } + } else { + for &seq in seqs { + let id = RequestId::new(20_000 + seq as u64); + let pr = ex + .execute_prefill(PrefillPlan { + requests: &[prefill_item(id, g.prompt(seq))], + }) + .expect("TP2 prefill"); + fold( + &mut stats, + seq, + 0, + &top_logprobs(pr.requests[0].first_token_logprob.as_ref()), + ); + for step in 0..g.decode_len { + let dr = ex + .execute_decode(DecodePlan { + requests: &[decode_item(id, g.decode(seq, step))], + }) + .expect("TP2 decode"); + fold( + &mut stats, + seq, + step + 1, + &top_logprobs(dr.requests[0].logprob.as_ref()), + ); + } + ex.drop_request(id).expect("TP2 drop request"); + } + } + (stats, fingerprint) +} + fn prompt_lens_label(g: &Golden) -> String { (0..g.num_seqs) .map(|seq| format!("{seq}:{}", g.prompt_len(seq))) @@ -658,6 +746,12 @@ fn build_executor(model_path: &str) -> Qwen35Executor { .expect("build Qwen3.5 logits executor") } +fn build_tp2_executor(model_path: &str) -> Qwen35TpExecutor { + let devices = common::tp2_device_ordinals(); + Qwen35TpExecutor::from_runtime_with_capacity(model_path, false, &devices, MAX_EXECUTOR_BATCH) + .expect("build Qwen3.5 TP2 logits executor") +} + #[test] fn pega_logprobs_match_hf_golden_within_qwen35_tolerance() { let Some(model_path) = model_path_or_skip() else { @@ -744,3 +838,57 @@ fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance() { "long sequential Qwen3.5 replay must reproduce identical logprobs" ); } + +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let Some(golden) = Golden::load_for(&model_path, false) else { + return; + }; + if !check_fixture_metadata(&model_path, &golden) { + return; + } + report_fixture_shape(&golden); + let all: Vec = (0..golden.num_seqs).collect(); + + let ex = build_tp2_executor(&model_path); + let (stats, fp1) = run_tp(&golden, &ex, &all, false); + report_and_assert("TP2 sequential eager", &stats); + let (_, fp2) = run_tp(&golden, &ex, &all, false); + assert_eq!( + fp1, fp2, + "TP2 sequential Qwen3.5 replay must reproduce identical logprobs" + ); + + let batched_n = all.len().min(MAX_EXECUTOR_BATCH); + let (batched, _) = run_tp(&golden, &ex, &all[..batched_n], true); + report_and_assert("TP2 batched eager", &batched); +} + +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance_tp2() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let Some(golden) = Golden::load_for(&model_path, true) else { + return; + }; + if !check_fixture_metadata(&model_path, &golden) { + return; + } + report_fixture_shape(&golden); + let all: Vec = (0..golden.num_seqs).collect(); + + let ex = build_tp2_executor(&model_path); + let (stats, fp1) = run_tp(&golden, &ex, &all, false); + report_and_assert("TP2 long sequential eager", &stats); + let (_, fp2) = run_tp(&golden, &ex, &all, false); + assert_eq!( + fp1, fp2, + "TP2 long sequential Qwen3.5 replay must reproduce identical logprobs" + ); +} diff --git a/openinfer-qwen35-4b/tests/serving_tp2.rs b/openinfer-qwen35-4b/tests/serving_tp2.rs new file mode 100644 index 000000000..724c13a34 --- /dev/null +++ b/openinfer-qwen35-4b/tests/serving_tp2.rs @@ -0,0 +1,334 @@ +use std::net::TcpListener; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::{Context, Result, anyhow, bail}; +use openinfer_core::engine::EngineLoadOptions; +use reqwest::Client; +use serde_json::{Value, json}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +mod common; + +const DEFAULT_MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); +const MODEL_NAME: &str = "qwen35-tp2-serving-smoke"; +const HTTP_TIMEOUT: Duration = Duration::from_secs(120); + +struct Qwen35Tp2Server { + base_url: String, + shutdown: CancellationToken, + task: JoinHandle>, +} + +impl Qwen35Tp2Server { + async fn shutdown(self) -> Result<()> { + self.shutdown.cancel(); + tokio::time::timeout(Duration::from_secs(30), self.task) + .await + .context("timed out waiting for Qwen3.5 TP2 frontend shutdown")? + .context("Qwen3.5 TP2 frontend task panicked")? + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires two CUDA devices, CUDA-12 NCCL, Qwen3.5 weights, and real HTTP frontend startup"] +async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { + let engine_model_path = PathBuf::from(get_model_path()); + let frontend_model_path = PathBuf::from(get_frontend_model_path(&engine_model_path)); + let server = spawn_ready_server(engine_model_path, frontend_model_path, 1).await?; + let client = test_client()?; + + assert_models_endpoint(&client, &server.base_url).await?; + assert_non_streaming_completion(&client, &server.base_url).await?; + assert_streaming_completion(&client, &server.base_url).await?; + assert_concurrent_completions(&client, &server.base_url).await?; + assert_invalid_cuda_graph_tp_startup_fails(&get_model_path())?; + + server.shutdown().await +} + +async fn spawn_ready_server( + engine_model_path: PathBuf, + frontend_model_path: PathBuf, + max_prefill_tokens: usize, +) -> Result { + let device_ordinals = common::tp2_device_ordinals(); + let handle = tokio::task::spawn_blocking(move || { + openinfer_qwen35_4b::start_engine_with_capacity( + &engine_model_path, + EngineLoadOptions { + enable_cuda_graph: false, + device_ordinals, + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + max_prefill_tokens, + ) + }) + .await + .context("Qwen3.5 TP2 engine loader thread panicked")??; + + let port = reserve_loopback_port()?; + let base_url = format!("http://127.0.0.1:{port}"); + let shutdown = CancellationToken::new(); + let server_shutdown = shutdown.clone(); + let mut task = tokio::spawn(async move { + openinfer_vllm_frontend::serve( + std::future::ready(Ok(handle)), + &frontend_model_path, + vec![MODEL_NAME.to_string()], + port, + None, + server_shutdown, + ) + .await + }); + + let client = test_client()?; + let health_result = tokio::select! { + result = wait_for_health(&client, &base_url) => result, + result = &mut task => { + return match result { + Ok(Ok(())) => Err(anyhow!("Qwen3.5 TP2 frontend exited before becoming healthy")), + Ok(Err(error)) => Err(error).context("Qwen3.5 TP2 frontend exited before becoming healthy"), + Err(error) => Err(error).context("Qwen3.5 TP2 frontend task panicked"), + }; + } + }; + + if let Err(error) = health_result { + shutdown.cancel(); + let _ = tokio::time::timeout(Duration::from_secs(30), task).await; + return Err(error).context("Qwen3.5 TP2 frontend did not become healthy"); + } + + Ok(Qwen35Tp2Server { + base_url, + shutdown, + task, + }) +} + +async fn assert_models_endpoint(client: &Client, base_url: &str) -> Result<()> { + let models: Value = client + .get(format!("{base_url}/v1/models")) + .send() + .await? + .error_for_status()? + .json() + .await?; + let advertised = models["data"] + .as_array() + .ok_or_else(|| anyhow!("/v1/models response has no data array: {models}"))?; + if !advertised.iter().any(|model| model["id"] == MODEL_NAME) { + bail!("/v1/models did not advertise {MODEL_NAME}: {models}"); + } + Ok(()) +} + +async fn assert_non_streaming_completion(client: &Client, base_url: &str) -> Result<()> { + let completion = post_completion(client, base_url, completion_body(false, 5, 1)).await?; + let choice = &completion["choices"][0]; + let text = choice["text"] + .as_str() + .ok_or_else(|| anyhow!("non-streaming completion has no text: {completion}"))?; + if text.is_empty() { + bail!("non-streaming completion returned empty text for max_tokens > 0"); + } + let finish_reason = choice["finish_reason"] + .as_str() + .ok_or_else(|| anyhow!("non-streaming completion has no finish_reason: {completion}"))?; + if finish_reason != "length" { + bail!("expected length finish_reason for ignore_eos request, got {completion}"); + } + assert_usage(&completion, 5)?; + assert_logprobs(&completion)?; + Ok(()) +} + +async fn assert_streaming_completion(client: &Client, base_url: &str) -> Result<()> { + let stream = post_completion_stream(client, base_url, completion_body(true, 4, 0)).await?; + let data_lines: Vec<&str> = stream + .lines() + .filter(|line| line.starts_with("data: ")) + .collect(); + if !data_lines.iter().any(|line| line.trim() == "data: [DONE]") { + bail!("streaming completion did not emit terminal data: [DONE]: {stream}"); + } + if !data_lines + .iter() + .filter(|line| line.trim() != "data: [DONE]") + .any(|line| line.contains("\"choices\"")) + { + bail!("streaming completion did not emit any choice payloads: {stream}"); + } + Ok(()) +} + +async fn assert_concurrent_completions(client: &Client, base_url: &str) -> Result<()> { + let first = post_completion(client, base_url, completion_body(false, 3, 0)); + let second = post_completion(client, base_url, alternate_completion_body(false, 3, 1)); + let (first, second) = tokio::try_join!(first, second)?; + assert_usage(&first, 3)?; + assert_usage(&second, 3)?; + assert_logprobs(&second)?; + Ok(()) +} + +fn assert_invalid_cuda_graph_tp_startup_fails(model_path: &str) -> Result<()> { + let error = match openinfer_qwen35_4b::start_engine_with_capacity( + Path::new(model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: common::tp2_device_ordinals(), + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + 1, + ) { + Ok(_) => bail!("TP2 + CUDA Graph must fail before serving requests"), + Err(error) => error, + }; + let message = error.to_string(); + if !message.contains("eager execution only") { + bail!("unexpected TP2 + CUDA Graph startup error: {message}"); + } + Ok(()) +} + +async fn post_completion(client: &Client, base_url: &str, body: Value) -> Result { + client + .post(format!("{base_url}/v1/completions")) + .json(&body) + .send() + .await? + .error_for_status()? + .json() + .await + .context("failed to parse non-streaming completion response") +} + +async fn post_completion_stream(client: &Client, base_url: &str, body: Value) -> Result { + client + .post(format!("{base_url}/v1/completions")) + .json(&body) + .send() + .await? + .error_for_status()? + .text() + .await + .context("failed to read streaming completion response") +} + +fn completion_body(stream: bool, max_tokens: usize, logprobs: usize) -> Value { + let mut body = json!({ + "model": MODEL_NAME, + "prompt": [151644, 872, 198, 9707, 151645, 198, 151644, 77091, 198], + "max_tokens": max_tokens, + "temperature": 0.0, + "ignore_eos": true, + "stream": stream + }); + if logprobs > 0 { + body["logprobs"] = json!(logprobs); + } + body +} + +fn alternate_completion_body(stream: bool, max_tokens: usize, logprobs: usize) -> Value { + let mut body = json!({ + "model": MODEL_NAME, + "prompt": [151644, 872, 198, 3838, 374, 220, 17, 489, 220, 17, 30, 151645, 198, 151644, 77091, 198], + "max_tokens": max_tokens, + "temperature": 0.0, + "ignore_eos": true, + "stream": stream + }); + if logprobs > 0 { + body["logprobs"] = json!(logprobs); + } + body +} + +fn assert_usage(completion: &Value, expected_completion_tokens: usize) -> Result<()> { + let completion_tokens = completion["usage"]["completion_tokens"] + .as_u64() + .ok_or_else(|| { + anyhow!("completion response has no usage.completion_tokens: {completion}") + })?; + if completion_tokens != expected_completion_tokens as u64 { + bail!( + "expected {expected_completion_tokens} completion tokens, got {completion_tokens}: {completion}" + ); + } + let prompt_tokens = completion["usage"]["prompt_tokens"] + .as_u64() + .ok_or_else(|| anyhow!("completion response has no usage.prompt_tokens: {completion}"))?; + if prompt_tokens == 0 { + bail!("completion response reported zero prompt tokens: {completion}"); + } + Ok(()) +} + +fn assert_logprobs(completion: &Value) -> Result<()> { + let logprobs = &completion["choices"][0]["logprobs"]; + if logprobs.is_null() { + bail!("completion requested logprobs but response has null logprobs: {completion}"); + } + let token_logprobs = logprobs["token_logprobs"] + .as_array() + .ok_or_else(|| anyhow!("logprobs.token_logprobs is not an array: {completion}"))?; + if token_logprobs.is_empty() { + bail!("logprobs.token_logprobs is empty: {completion}"); + } + if !token_logprobs.iter().all(|value| value.as_f64().is_some()) { + bail!("logprobs.token_logprobs contains non-finite values: {completion}"); + } + Ok(()) +} + +async fn wait_for_health(client: &Client, base_url: &str) -> Result<()> { + let health_url = format!("{base_url}/health"); + let deadline = tokio::time::Instant::now() + Duration::from_secs(180); + loop { + if tokio::time::Instant::now() >= deadline { + bail!("timed out waiting for Qwen3.5 TP2 frontend health at {health_url}"); + } + + match client + .get(&health_url) + .timeout(Duration::from_secs(2)) + .send() + .await + { + Ok(response) if response.status().is_success() => return Ok(()), + Ok(_) | Err(_) => tokio::time::sleep(Duration::from_millis(500)).await, + } + } +} + +fn test_client() -> Result { + Client::builder() + .no_proxy() + .timeout(HTTP_TIMEOUT) + .build() + .context("failed to build HTTP test client") +} + +fn reserve_loopback_port() -> Result { + let listener = TcpListener::bind(("127.0.0.1", 0)) + .context("failed to reserve loopback port for Qwen3.5 TP2 serving test")?; + Ok(listener.local_addr()?.port()) +} + +fn get_model_path() -> String { + std::env::var("OPENINFER_TEST_MODEL_PATH").unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string()) +} + +fn get_frontend_model_path(engine_model_path: &Path) -> String { + std::env::var("OPENINFER_TEST_FRONTEND_MODEL_PATH") + .unwrap_or_else(|_| engine_model_path.to_string_lossy().into_owned()) +} diff --git a/openinfer-server/src/config.rs b/openinfer-server/src/config.rs index e478c044f..7294adabf 100644 --- a/openinfer-server/src/config.rs +++ b/openinfer-server/src/config.rs @@ -364,6 +364,7 @@ fn consumed_args(model_type: ModelType) -> &'static [&'static str] { #[cfg(feature = "qwen35-4b")] ModelType::Qwen35 => &[ "device_ordinal", + "tp_size", "cuda_graph", "max_prefill_tokens", "max_batch", @@ -629,7 +630,7 @@ fn parse_lora_module_fields(name: &str, path: &str) -> Result (Args, BTreeSet) { use clap::FromArgMatches; let matches = Args::command() @@ -674,6 +675,15 @@ mod tests { } } + #[cfg(feature = "qwen35-4b")] + #[test] + fn qwen35_accepts_tp_size() { + let (args, provided) = + parse_with_provided(&["openinfer", "--tp-size", "2", "--cuda-graph=false"]); + args.validate(ModelType::Qwen35, &provided) + .expect("Qwen3.5 should accept --tp-size for eager TP startup"); + } + #[test] fn parses_lora_modules_name_equals_path() { assert_eq!( diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index 40a3de810..8c276837f 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -303,18 +303,19 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result openinfer_qwen35_4b::start_engine( + ModelType::Qwen35 => openinfer_qwen35_4b::launch_with_options( &args.model_path, - EngineLoadOptions { - enable_cuda_graph: args.cuda_graph, - device_ordinals: vec![args.device_ordinal], - seed: 42, - ..EngineLoadOptions::default() + openinfer_qwen35_4b::Qwen35LaunchOptions { + device_ordinal: args.device_ordinal, + tp_size: args.tp_size, + cuda_graph: args.cuda_graph, + max_batch: args + .max_batch + .unwrap_or(openinfer_qwen35_4b::runtime::MAX_BATCH), + max_prefill_tokens: args + .max_prefill_tokens + .unwrap_or(openinfer_qwen35_4b::DEFAULT_MAX_PREFILL_TOKENS), }, - args.max_batch - .unwrap_or(openinfer_qwen35_4b::runtime::MAX_BATCH), - args.max_prefill_tokens - .unwrap_or(openinfer_qwen35_4b::DEFAULT_MAX_PREFILL_TOKENS), ) .context("failed to start Qwen3.5 engine")?, };