Feat/qwen35 tp phase1#627
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b380a1baec
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4f50a0f01
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| &format!("CUmodule {module_name}[16] = {{0}};"), | ||
| ) | ||
| .replace( | ||
| &format!("CUfunction {function_name} = NULL;"), | ||
| &format!("CUfunction {function_name}[16] = {{0}};"), |
There was a problem hiding this comment.
When Qwen3.5 is launched on CUDA ordinal 16 or higher, the generated stubs index these fixed 16-entry arrays with the dev returned by cuCtxGetDevice in the patched load/launch paths, so module_name[dev] or function_name[dev] writes past the table and can corrupt memory. Please size the tables from the actual CUDA device count or fail closed when dev is outside the allocated range.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This issue was not caused by switching from a single handle to a per-device list. The real fix was to support multiple CUDA devices with separate handles. We kept the fixed 16-entry table, and added a check that rejects any device ordinal greater than 15 before indexing it. If reviewers find this approach suboptimal, please let me know, and I will update it accordingly.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4f50a0f01
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
hi thanks your contribution, please fix the conflicts |
a624e29 to
920ec4a
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 920ec4a9f3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a33275755f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 07dc15be4f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| }) | ||
| .expect("failed to spawn Qwen3.5 TP scheduler thread"); | ||
|
|
||
| Ok(SchedulerHandle::new_with_join_handle(submit_tx, join_handle).with_servable_len(servable)) |
There was a problem hiding this comment.
When Qwen3.5 starts with tp_size > 1, this returns an EngineHandle without with_kv_capacity, even though the TP backend already knows capacity_pages_for_requests() and page_size(). The vLLM bridge treats handle.kv_capacity() == None as num_gpu_blocks = 0 and omits cache size/concurrency in its ready response (openinfer-vllm-frontend/src/bridge.rs:65-80), so TP HTTP serving advertises zero KV blocks to the frontend/tooling despite having a real cache. Mirror the single-GPU path by attaching a KvCapacity built from the TP backend's page count and block size.
Useful? React with 👍 / 👎.
|
All conflicts fixed. |
| let mut result = None; | ||
| for (rank, recv) in pending.into_iter().enumerate() { | ||
| match recv | ||
| .recv() |
There was a problem hiding this comment.
A worker can return a rank-local CUDA/OOM error before its peers complete the corresponding NCCL collective. The controller then waits indefinitely on these unbounded recv() calls, and Drop can also block forever while joining a worker stuck in NCCL. Please add a runtime failure contract that aborts the communicator or process, marks the executor as poisoned, and prevents unbounded cleanup.
There was a problem hiding this comment.
Solved. The runtime fix adds a shared poison state, propagates worker errors and panics, uses bounded response and shutdown waits, and aborts the process if NCCL workers cannot exit, preventing indefinite runtime and cleanup hangs.
| request_id, | ||
| phase: TpRequestPhase::Prefilling, | ||
| kv: self.model.alloc_kv(), | ||
| recurrent: RecurrentState::new(self.model.device_ctx(), self.model.config())?, |
There was a problem hiding this comment.
Each new TP request allocates a full replicated recurrent and convolution state on every rank. For Qwen3.5-4B this is about 49.125 MiB per request per rank, so the default max_batch=64 requires about 3.07 GiB per rank after the KV pool has already consumed 85% of the available memory. Please reserve or preallocate this state before sizing the KV pool, or derive and advertise a lower capacity from the minimum available memory across ranks.
There was a problem hiding this comment.
Solved. The capacity fix calculates per-request recurrent-state memory, reserves GPU runtime headroom, derives a safe max_batch for each rank, uses the minimum across all ranks, and rejects startup if even one request cannot fit.
| .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 max_batch = effective_recurrent_capacity( |
There was a problem hiding this comment.
This preserves only 512 MiB, but a legal 20k-token TP2 prefill needs about 2.66 GiB of scratch. Recurrent states can consume that headroom under high occupancy and cause OOM. Please subtract the configured prefill scratch before calculating capacity, or cap max_prefill_tokens.
There was a problem hiding this comment.
Solved. Before calculating recurrent-state capacity, we subtract both the runtime reserve and the estimated prefill scratch allocation from the available memory. The scratch allocation is estimated using min(max_prefill_tokens, PREFILL_CHUNK_LEN), where PREFILL_CHUNK_LEN is currently 20,000 tokens. Thus, we reserve scratch for the smaller of the configured maximum prefill budget and the hard cap of 20k tokens.
Note that the subtracted value is estimated value calculated by GdrChunkwiseScratch35::estimate_bytes(...), this will be slightly bigger than 2.66 GiB.
5d84063 to
879044a
Compare
Squashed PR #627 from Mrtroll486/feat/qwen35-tp-phase1.
2 participants
Related
Background
This PR lands Qwen3.5-4B TP Phase 1.
Phase 1 is a correctness-first eager dense-TP milestone: it brings up TP2 dense full-attention / MLP sharding, worker/scheduler lifecycle, and the real serving path. This PR does not change the design contract doc
docs/models/qwen35/tp-design.md; if review finds that additional contract wording is needed, I will update the design-doc branch / PR #450 rather than patching the design contract directly in this implementation PR.Phase 2 design work, including TP unified steps, TP CUDA Graph, and sharded linear-attention/GDR state, should be decided in follow-up RFC issue discussion. Any unsettled design points mentioned below are follow-ups, not Phase 1 design commitments.
What Changed
Qwen3.5 TP runtime
rank/world_sizeTP > 1 && CUDA Graphq/k/v/ogate/up/downlm_headDropRequestclears rank-local request stateQwen3.5 test framework
OPENINFER_TEST_TP_DEVICES.0,1.1,2or2,3.OPENINFER_TEST_MODEL_PATHselects the real Qwen3.5 weights path.OPENINFER_TEST_FRONTEND_MODEL_PATHoptionally selects frontend tokenizer/config metadata for HTTP serving tests.OPENINFER_TEST_FRONTEND_MODEL_PATHfalls back toOPENINFER_TEST_MODEL_PATHwhen unset.Qwen35TpExecutorTP > 1 && CUDA Graphopeninfer_vllm_frontend::serve/v1/modelsand/v1/completionsServer / frontend launch path
tp_sizedevice_ordinalcuda_graphmax_prefill_tokensopeninfer-serverQwen3.5 startup to use the new launch options.tp_size > 1maps to devices0..tp_sizein the server path.bench_servingwith--max-prefill-tokensfor Qwen3/Qwen3.5 chunked-prefill control.openinfer_vllm_frontendcrate.Docs
docs/models/qwen35/tp-implementation.md.tp-design.md.Testing
TP2 correctness
108positions, mean0.0258, p990.0801, max0.129872positions, mean0.0257, p990.0809, max0.12984097/819218positions, mean0.0232, p990.0792, max0.1035Scheduler / serving
/v1/models/v1/completions/v1/completionsmax_prefill_tokens=1TP2 + CUDA GraphTP1 regression
Compile / focused tests
serving_tp2test target compiles.openinfer-server --features qwen35-4btest target compiles.Out Of Scope
RunUnifiedStepmixed prefill+decode executionlm_headReviewer Notes / Open Design Points
1. TP unified step
The Phase 1 TP scheduler currently selects eager prefill-only or eager decode-only steps. It does not select a mixed prefill+decode unified step under TP.
This PR’s position is that Phase 1 should first prove dense TP correctness, worker-owned state, and rank-local cleanup/reuse. The semantics, scheduling policy, acceptance gates, and sequencing for TP
RunUnifiedStepshould be decided in a follow-up RFC issue.If reviewers want
tp-design.mdto explicitly say that Phase 1 does not coverRunUnifiedStep, I will update the design-doc branch / PR #450 rather than changing the design contract directly in this PR.2. Slot / slot-compaction wording
There are two concepts that are easy to conflate:
BatchDecodeGraphStateTP2 Phase 1 is eager-only and does not use TP CUDA Graph recurrent slots. It validates consistent rank-local request cleanup across finish/drop/cancel/reuse, which is not the same thing as TP CUDA Graph slot compaction.
If this distinction should become part of the design contract, I will update PR #450 / the design-doc branch.
3. Gated
q_projphysical layoutThe implementation confirmed that Qwen3.5 full-attention
q_projuses a head-interleaved physical layout:[head q][head gate]. TP sharding must preserve that layout. Rebuilding local shards as[all local q][all local gate]breaks numerics.This is recorded in the implementation doc for now. If reviewers consider it part of the long-term partition contract, I will add the generic formula and TP2 example to PR #450 / the design-doc branch.
Docs
Added:
docs/models/qwen35/tp-implementation.mdThis doc records the Phase 1 implementation and validation evidence. The stable design contract remains in
docs/models/qwen35/tp-design.md/ PR #450.Type of Change
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Documentation update
Checklist
My code follows the style guidelines of this project (see
docs/conventions/coding-style.md).I have performed a self-review of my own code.
I have formatted my commits according to Commitizen conventions.
I have run the local test suite and all tests pass (see
CLAUDE.md).