diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5ee080c1f..149cb58eb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -319,3 +319,71 @@ jobs: CUDA_PATH: /usr/local/cuda PEGAINFER_CUDA_SM: "80" PEGAINFER_NVCC_JOBS: "2" + + qwen35-cuda: + if: github.event_name == 'push' || github.event.pull_request.draft == false + name: Qwen3.5 CUDA ${{ matrix.gate }} (sm_80) + runs-on: ubuntu-latest + timeout-minutes: 45 + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + strategy: + fail-fast: false + matrix: + include: + - gate: Compile + cargo_command: check + trailing_args: "" + - gate: Clippy + cargo_command: clippy + trailing_args: "-- -D warnings" + + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + persist-credentials: false + submodules: recursive + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@v1 + with: + toolchain: nightly-2026-07-10 + components: clippy + + - name: Setup sccache + uses: mozilla-actions/sccache-action@v0.0.10 + with: + version: "v0.16.0" + + - name: Install CUDA toolkit + uses: Jimver/cuda-toolkit@v0.2.35 + with: + cuda: "13.0.2" + linux-local-args: '["--toolkit"]' + method: network + sub-packages: '["nvcc", "nvrtc-dev", "cudart-dev"]' + non-cuda-sub-packages: '["libcublas-dev", "libcurand-dev"]' + + - name: Install Python + uses: actions/setup-python@v6 + with: + python-version: "3.10" + + - name: Install Qwen3.5 build dependencies + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler libibverbs-dev + python -m pip install "triton==3.7.1" + + - name: Run Qwen3.5 ${{ matrix.gate }} + run: >- + cargo ${{ matrix.cargo_command }} --release --locked + -p pegainfer-qwen35 --features qwen35 --all-targets + ${{ matrix.trailing_args }} + env: + CUDA_PATH: /usr/local/cuda + PEGAINFER_CUDA_SM: "80" + PEGAINFER_NVCC_JOBS: "2" + PEGAINFER_TRITON_PYTHON: python diff --git a/docs/benchmarks/qwen35-tp2-phase2a-multiturn.json b/docs/benchmarks/qwen35-tp2-phase2a-multiturn.json new file mode 100644 index 000000000..3c58b57ef --- /dev/null +++ b/docs/benchmarks/qwen35-tp2-phase2a-multiturn.json @@ -0,0 +1 @@ +{"date":"2026-08-19T14:24:36Z","mode":"multi_turn","backend":"openai-chat","model_id":"qwen35-tp2-phase2a","num_conversations":12,"max_model_len":2048,"turns_per_conversation":3,"multi_turn_concurrency":4,"inter_turn_delay_ms":0,"server_commit":"8829189992d9290f1aa128fb00950dd854f43211","client_commit":"2b7fcbf52782f8729fd6ce6c9ab803617d72897b","gpu":"2x_RTX_3090_sm86","tp_size":"2","server_max_batch":"2","server_max_prefill_tokens":"64","duration":13.864881841999999,"conversations_completed":12,"conversations_failed":0,"avg_turns_completed":3.6666666666666665,"avg_conversation_duration_ms":10451.412879833333,"completed":44,"failed":0,"total_input_tokens":12496,"total_output_tokens":1056,"request_throughput":3.173485392909272,"input_throughput":901.2698515862332,"output_throughput":76.16364942982253,"total_token_throughput":977.4335010160559,"mean_ttft_ms":727.5413035227273,"median_ttft_ms":727.6886789999999,"std_ttft_ms":146.66137831986737,"p50_ttft_ms":727.6886789999999,"p90_ttft_ms":847.6425581,"p99_ttft_ms":899.2532077699999,"mean_tpot_ms":21.70376772332016,"median_tpot_ms":21.55133260869565,"std_tpot_ms":0.8264917019729615,"p50_tpot_ms":21.55133260869565,"p90_tpot_ms":22.558429852173912,"p99_tpot_ms":22.933247999565218,"mean_itl_ms":19.96713433909089,"median_itl_ms":20.5969695,"std_itl_ms":6.481995568729013,"p50_itl_ms":20.5969695,"p90_itl_ms":26.920911500000003,"p99_itl_ms":30.078894709999997,"mean_e2el_ms":1226.7279611590914,"median_e2el_ms":1222.6157295,"std_e2el_ms":153.48850861286388,"p50_e2el_ms":1222.6157295,"p90_e2el_ms":1365.4634053,"p99_e2el_ms":1406.7637406399997,"per_turn_metrics":[{"turn_index":0,"num_samples":12,"completed":12,"failed":0,"total_input_tokens":1824,"total_output_tokens":288,"request_throughput":0.8654960162479833,"input_throughput":131.55539446969345,"output_throughput":20.7719043899516,"total_token_throughput":152.32729885964505,"mean_ttft_ms":583.7918206666667,"median_ttft_ms":657.463881,"std_ttft_ms":198.93882849514316,"p50_ttft_ms":657.463881,"p90_ttft_ms":691.2390482,"p99_ttft_ms":742.9091798999999,"mean_tpot_ms":21.169481184782605,"median_tpot_ms":21.11208156521739,"std_tpot_ms":0.1470605595878701,"p50_tpot_ms":21.11208156521739,"p90_tpot_ms":21.43137102173913,"p99_tpot_ms":21.507665196521735,"mean_itl_ms":19.47549855000002,"median_itl_ms":20.379788,"std_itl_ms":6.2034251783074525,"p50_itl_ms":20.379788,"p90_itl_ms":21.269388,"p99_itl_ms":30.728750939999998,"mean_e2el_ms":1070.6898879166667,"median_e2el_ms":1147.035556,"std_e2el_ms":199.1779687190421,"p50_e2el_ms":1147.035556,"p90_e2el_ms":1175.5356078000002,"p99_e2el_ms":1228.5809922},{"turn_index":1,"num_samples":12,"completed":12,"failed":0,"total_input_tokens":2880,"total_output_tokens":288,"request_throughput":0.8654960162479833,"input_throughput":207.719043899516,"output_throughput":20.7719043899516,"total_token_throughput":228.49094828946758,"mean_ttft_ms":714.34124925,"median_ttft_ms":720.4639175000001,"std_ttft_ms":16.123405097003147,"p50_ttft_ms":720.4639175000001,"p90_ttft_ms":722.1488714999999,"p99_ttft_ms":732.58807688,"mean_tpot_ms":21.50762221376812,"median_tpot_ms":21.512210391304347,"std_tpot_ms":0.05180367387233731,"p50_tpot_ms":21.512210391304347,"p90_tpot_ms":21.557149665217388,"p99_tpot_ms":21.592176384782608,"mean_itl_ms":19.786716913333336,"median_itl_ms":20.530554,"std_itl_ms":6.297680823403235,"p50_itl_ms":20.530554,"p90_itl_ms":26.572883,"p99_itl_ms":30.4137518,"mean_e2el_ms":1209.0165601666667,"median_e2el_ms":1214.526274,"std_e2el_ms":15.950441454258286,"p50_e2el_ms":1214.526274,"p90_e2el_ms":1216.6144174,"p99_e2el_ms":1227.28499475},{"turn_index":2,"num_samples":9,"completed":9,"failed":0,"total_input_tokens":2952,"total_output_tokens":216,"request_throughput":0.6491220121859875,"input_throughput":212.91201999700388,"output_throughput":15.578928292463699,"total_token_throughput":228.49094828946758,"mean_ttft_ms":778.062378,"median_ttft_ms":783.74919,"std_ttft_ms":19.224187823957255,"p50_ttft_ms":783.74919,"p90_ttft_ms":789.8680276000001,"p99_ttft_ms":805.37482096,"mean_tpot_ms":22.044310932367154,"median_tpot_ms":21.928995999999998,"std_tpot_ms":0.21255355180988741,"p50_tpot_ms":21.928995999999998,"p90_tpot_ms":22.314300469565218,"p99_tpot_ms":22.48564071652174,"mean_itl_ms":20.280475453333345,"median_itl_ms":20.668806,"std_itl_ms":6.545678930929194,"p50_itl_ms":20.668806,"p90_itl_ms":26.9185278,"p99_itl_ms":29.846178479999995,"mean_e2el_ms":1285.0815294444446,"median_e2el_ms":1287.6416179999999,"std_e2el_ms":17.447565570324695,"p50_e2el_ms":1287.6416179999999,"p90_e2el_ms":1303.2953988,"p99_e2el_ms":1310.6299300800001},{"turn_index":3,"num_samples":8,"completed":8,"failed":0,"total_input_tokens":3328,"total_output_tokens":192,"request_throughput":0.5769973441653222,"input_throughput":240.03089517277405,"output_throughput":13.847936259967733,"total_token_throughput":253.87883143274178,"mean_ttft_ms":845.3753072500001,"median_ttft_ms":846.2088815,"std_ttft_ms":13.52094063421677,"p50_ttft_ms":846.2088815,"p90_ttft_ms":855.3806129000001,"p99_ttft_ms":871.3374689899999,"mean_tpot_ms":22.59411058695652,"median_tpot_ms":22.542907565217387,"std_tpot_ms":0.1418984913926098,"p50_tpot_ms":22.542907565217387,"p90_tpot_ms":22.768807678260867,"p99_tpot_ms":22.91283586782609,"mean_itl_ms":20.786269600000008,"median_itl_ms":20.8269005,"std_itl_ms":6.734698925762719,"p50_itl_ms":20.8269005,"p90_itl_ms":27.4097134,"p99_itl_ms":29.42091591,"mean_e2el_ms":1365.03985075,"median_e2el_ms":1364.992177,"std_e2el_ms":14.886185612840757,"p50_e2el_ms":1364.992177,"p90_e2el_ms":1379.5870000000002,"p99_e2el_ms":1393.6527202},{"turn_index":4,"num_samples":3,"completed":3,"failed":0,"total_input_tokens":1512,"total_output_tokens":72,"request_throughput":0.21637400406199583,"input_throughput":109.05249804724589,"output_throughput":5.1929760974879,"total_token_throughput":114.24547414473379,"mean_ttft_ms":889.5522186666667,"median_ttft_ms":883.5537770000001,"std_ttft_ms":15.724934853616045,"p50_ttft_ms":883.5537770000001,"p90_ttft_ms":905.5880658,"p99_ttft_ms":910.54578078,"mean_tpot_ms":21.229618652173915,"median_tpot_ms":22.913742956521734,"std_tpot_ms":2.3978735945256315,"p50_tpot_ms":22.913742956521734,"p90_tpot_ms":22.932007860869565,"p99_tpot_ms":22.936117464347824,"mean_itl_ms":19.530963160000006,"median_itl_ms":20.905244,"std_itl_ms":7.1278942722312735,"p50_itl_ms":20.905244,"p90_itl_ms":27.1184224,"p99_itl_ms":28.359043100000008,"mean_e2el_ms":1377.8334476666664,"median_e2el_ms":1401.022329,"std_e2el_ms":40.12772579823275,"p50_e2el_ms":1401.022329,"p90_e2el_ms":1409.0804506,"p99_e2el_ms":1410.8935279599998}]} diff --git a/docs/benchmarks/qwen35-tp2-phase2a-multiturn.md b/docs/benchmarks/qwen35-tp2-phase2a-multiturn.md new file mode 100644 index 000000000..b96d1b9ff --- /dev/null +++ b/docs/benchmarks/qwen35-tp2-phase2a-multiturn.md @@ -0,0 +1,256 @@ +# Qwen3.5 Phase 2A TP2 Multi-Turn Serving Gate + +> **TL;DR:** A real 2x RTX 3090 Qwen3.5 TP2 server completed 12/12 dependent conversations and 44/44 measured turns at client concurrency 4 over server capacity 2, then admitted and completed another 4/4 conversations and 8/8 turns without restart; all requests produced the configured token counts, history grew on every turn, shutdown was clean, and both GPUs returned to idle. + +This is a correctness and lifecycle gate for Phase 2A. It is not a comparative +performance claim. The run was captured on 2026-08-19 UTC. + +## Pass Criteria + +The run passes only when all of the following hold: + +1. The dry run builds 12 conversations with 2-5 turns and no server traffic. +2. The primary run completes 12/12 conversations and 44/44 measured turns with zero failures. +3. Every measured turn returns exactly 24 output tokens, for 1,056 output tokens total. +4. Per-turn sample counts prove mixed conversation lengths, and prompt accounting grows with carried chat history. +5. Client concurrency 4 completes against `max_batch=2` without deadlock or rank divergence. +6. Without restarting the server, a second probe completes 4/4 newly admitted conversations and 8/8 turns with zero failures. +7. Graceful shutdown exits the scheduler and releases both GPUs. + +All seven criteria passed. + +## Reproduction Pins + +| Component | Value | +| --- | --- | +| PegaInfer server commit | `8829189992d9290f1aa128fb00950dd854f43211` | +| PegaInfer binary SHA-256 | `97ab83bbf666add4e4b0854c1bea77688279ca27dc87bc15641fb7d7e29cd9af` | +| vLLM client commit | `2b7fcbf52782f8729fd6ce6c9ab803617d72897b` | +| `vllm-bench` version | `0.1.0` | +| `vllm-bench` binary SHA-256 | `8131ed513d22da21186eb7ccba06dfb6d0c8657624bb72a733948841a5dd1ffe` | +| GPUs | 2x NVIDIA RTX 3090, SM86, 24 GiB each | +| Model fixture | [`Qwen/Qwen3.5-4B`](https://huggingface.co/Qwen/Qwen3.5-4B/tree/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a), revision `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` | +| `config.json` SHA-256 | `ddc63e1c717afa86c865bb5e01313d89d72bb53b97ad4a8a03ba8510c0621670` | +| `tokenizer.json` SHA-256 | `5f9e4d4901a92b997e463c1f46055088b6cca5ca61a6522d1b9f64c4bb81cb42` | +| `model.safetensors.index.json` SHA-256 | `cf3f798ee02ba45f9622aa8892a47369ab667d0afbf154ee7c2212de42e6302d` | +| `model.safetensors-00001-of-00002.safetensors` SHA-256 | `26a93f066e1916adb13453dae5a0c707c0fbc71299ed98779571a907b8e74c61` | +| `model.safetensors-00002-of-00002.safetensors` SHA-256 | `cb544bd9bfae93dc59b0f22b292f5933573854a7f9b97835c67060d7d910e188` | +| TP/runtime | TP=2, eager, CUDA Graph disabled, `max_batch=2`, `max_prefill_tokens=64`, scheduler policy `off` | + +The repository's older pinned vLLM revision predates the Rust `vllm-bench` +workspace. The client was therefore built from the exact maintained upstream +commit above. Source inspection confirmed that `--multi-turn` runs turns +sequentially within each conversation, appends each assistant response, and +sends the accumulated message history on the next turn. A separate sequential +chat driver was not needed. + +Set these paths before running the commands. `MODEL_DIR` must point to the +pinned public revision above and match the complete config, tokenizer, index, +and weight-shard hash manifest; no private absolute model path is part of the +reproduction contract. + +```bash +MODEL_DIR=/path/to/public/Qwen3.5-4B +VLLM_SRC=/tmp/vllm-phase2a-bench-src +VLLM_BENCH="$VLLM_SRC/rust/target/release/vllm-bench" +RESULT_DIR=/tmp/qwen35-tp2-phase2a-results +``` + +Build the pinned client: + +```bash +git clone --filter=blob:none --no-checkout https://github.com/vllm-project/vllm.git "$VLLM_SRC" +git -C "$VLLM_SRC" fetch --depth 1 origin 2b7fcbf52782f8729fd6ce6c9ab803617d72897b +git -C "$VLLM_SRC" checkout --detach 2b7fcbf52782f8729fd6ce6c9ab803617d72897b +cargo build --release -p vllm-bench --manifest-path "$VLLM_SRC/rust/Cargo.toml" +"$VLLM_BENCH" --version +``` + +Build and start PegaInfer from the server commit above: + +```bash +PEGAINFER_CUDA_SM=86 \ +PEGAINFER_TRITON_PYTHON="$PWD/.venv/bin/python" \ +PROTOC="$PWD/.venv/protoc-root/usr/bin/protoc" \ +PROTOC_INCLUDE="$PWD/.venv/protoc-root/usr/include" \ +LD_LIBRARY_PATH="$PWD/.venv/protoc-root/usr/lib/x86_64-linux-gnu" \ +cargo build --offline --release --locked \ + -p pegainfer-server --no-default-features --features qwen35 +``` + +```bash +CUDA_VISIBLE_DEVICES=0,1 \ +PEGAINFER_CUDA_SM=86 \ +PEGAINFER_TRITON_PYTHON="$PWD/.venv/bin/python" \ +LD_LIBRARY_PATH="$PWD/.venv/lib/python3.10/site-packages/nvidia/nccl/lib" \ +RUST_LOG=info \ +target/release/pegainfer \ + --model-path "$MODEL_DIR" \ + --served-model-name qwen35-tp2-phase2a \ + --port 18080 \ + --tp-size 2 \ + --cuda-graph=false \ + --max-batch 2 \ + --max-prefill-tokens 64 \ + --qwen35-scheduler-policy off +``` + +## Workload Commands + +The dry run fixes the generated dataset before server traffic: + +```bash +"$VLLM_BENCH" \ + --backend openai-chat \ + --model qwen35-tp2-phase2a \ + --tokenizer "$MODEL_DIR" \ + --dataset-name random \ + --multi-turn \ + --multi-turn-min-turns 2 \ + --multi-turn-max-turns 5 \ + --random-input-len 128 \ + --per-turn-input-len 64 \ + --random-output-len 24 \ + --num-prompts 12 \ + --multi-turn-concurrency 4 \ + --max-model-len 2048 \ + --seed 446 \ + --dry-run +``` + +It produced 12 conversations, 44 turns, and 3,584 user-message tokens. +The measured primary run used the same seed and shape: + +```bash +mkdir -p "$RESULT_DIR" +"$VLLM_BENCH" \ + --backend openai-chat \ + --base-url http://127.0.0.1:18080 \ + --model qwen35-tp2-phase2a \ + --tokenizer "$MODEL_DIR" \ + --dataset-name random \ + --multi-turn \ + --multi-turn-min-turns 2 \ + --multi-turn-max-turns 5 \ + --random-input-len 128 \ + --per-turn-input-len 64 \ + --random-output-len 24 \ + --num-prompts 12 \ + --multi-turn-concurrency 4 \ + --max-model-len 2048 \ + --seed 446 \ + --ignore-eos \ + --temperature 0 \ + --ready-check-timeout-sec 30 \ + --percentile-metrics ttft,tpot,itl,e2el \ + --metric-percentiles 50,90,99 \ + --save-result \ + --save-detailed \ + --result-dir "$RESULT_DIR" \ + --result-filename qwen35-tp2-phase2a-multiturn.json \ + --metadata \ + server_commit=8829189992d9290f1aa128fb00950dd854f43211 \ + client_commit=2b7fcbf52782f8729fd6ce6c9ab803617d72897b \ + gpu=2x_RTX_3090_sm86 \ + tp_size=2 \ + server_max_batch=2 \ + server_max_prefill_tokens=64 +``` + +The post-workload probe was run immediately afterward, without restarting or +reloading the server: + +```bash +"$VLLM_BENCH" \ + --backend openai-chat \ + --base-url http://127.0.0.1:18080 \ + --model qwen35-tp2-phase2a \ + --tokenizer "$MODEL_DIR" \ + --dataset-name random \ + --multi-turn \ + --multi-turn-num-turns 2 \ + --random-input-len 32 \ + --per-turn-input-len 16 \ + --random-output-len 8 \ + --num-prompts 4 \ + --multi-turn-concurrency 4 \ + --max-model-len 512 \ + --seed 447 \ + --ignore-eos \ + --temperature 0 \ + --ready-check-timeout-sec 30 \ + --percentile-metrics ttft,tpot,itl,e2el \ + --metric-percentiles 50,90,99 \ + --save-result \ + --result-dir "$RESULT_DIR" \ + --result-filename qwen35-tp2-phase2a-post-readmission.json \ + --metadata \ + purpose=post_workload_readmission \ + server_commit=8829189992d9290f1aa128fb00950dd854f43211 \ + client_commit=2b7fcbf52782f8729fd6ce6c9ab803617d72897b \ + tp_size=2 \ + server_max_batch=2 +``` + +## Results + +| Metric | Primary run | Post-workload probe | +| --- | ---: | ---: | +| Conversations completed/failed | 12 / 0 | 4 / 0 | +| Measured turns completed/failed | 44 / 0 | 8 / 0 | +| Readiness requests, unmeasured | 1 | 1 | +| Client concurrency | 4 | 4 | +| Duration | 13.865 s | 0.786 s | +| Input tokens | 12,496 | 416 | +| Output tokens | 1,056 | 64 | +| Request throughput | 3.173 req/s | 10.185 req/s | +| Output throughput | 76.164 tok/s | 81.475 tok/s | + +Primary latency percentiles: + +| Metric | Mean | p50 | p90 | p99 | +| --- | ---: | ---: | ---: | ---: | +| TTFT | 727.54 ms | 727.69 ms | 847.64 ms | 899.25 ms | +| TPOT | 21.70 ms | 21.55 ms | 22.56 ms | 22.93 ms | +| ITL | 19.97 ms | 20.60 ms | 26.92 ms | 30.08 ms | +| E2EL | 1226.73 ms | 1222.62 ms | 1365.46 ms | 1406.76 ms | + +Turn-level evidence from the primary run: + +| Turn | Conversations reaching turn | Client-accounted input per request | Server prompt-token range | Output per request | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 12 | 152 | 138-140 | 24 | +| 2 | 12 | 240 | 212-214 | 24 | +| 3 | 9 | 328 | 286-288 | 24 | +| 4 | 8 | 416 | 360-362 | 24 | +| 5 | 3 | 504 | 435 | 24 | + +The `[12, 12, 9, 8, 3]` sample vector proves that conversations had mixed 2-5 +turn lengths. Both the client accounting and the server's rendered prompt +tokens increase per turn, which proves that later requests carried earlier +user and assistant messages instead of flattening the dataset to its first +turn. Client and server token counts differ because the server applies the chat +template and tokenizer to the accumulated messages. + +This client revision still writes the default value `turns_per_conversation=3` +in the top-level JSON even when min/max turn sampling is enabled. The measured +per-turn sample vector and `avg_turns_completed=3.6667` are the authoritative +values for this variable-length run. + +The client emitted a warning that `--ignore-eos` can interact with multi-turn +output limits. In this run, every one of the 44 server responses logged +`output_tokens=24` with `finish_reason=length`; the saved result totals exactly +`44 * 24 = 1,056` output tokens. The post-workload probe similarly totals +`8 * 8 = 64`. + +The primary result is stored in +[qwen35-tp2-phase2a-multiturn.json](qwen35-tp2-phase2a-multiturn.json), SHA-256 +`9247a35bc3f5475ec59397ddd2dd3461b9c033e1d27ef25f64c7517da5398676`. +The no-restart readmission result is stored in +[qwen35-tp2-phase2a-post-readmission.json](qwen35-tp2-phase2a-post-readmission.json), +SHA-256 `a60890446d3bb08f83da37bf7c14c2a9063eccbdc480ae3cbb028ccd52ce2573`. + +After the second run, Ctrl-C produced `scheduler: all handles dropped, exiting` +and process exit code 0. A subsequent `nvidia-smi` showed both TP devices at 1 +MiB used, 24,126 MiB free, and 0% utilization, confirming process-level resource +release. diff --git a/docs/benchmarks/qwen35-tp2-phase2a-post-readmission.json b/docs/benchmarks/qwen35-tp2-phase2a-post-readmission.json new file mode 100644 index 000000000..c28058bf5 --- /dev/null +++ b/docs/benchmarks/qwen35-tp2-phase2a-post-readmission.json @@ -0,0 +1 @@ +{"date":"2026-08-19T14:25:00Z","mode":"multi_turn","backend":"openai-chat","model_id":"qwen35-tp2-phase2a","num_conversations":4,"max_model_len":512,"turns_per_conversation":2,"multi_turn_concurrency":4,"inter_turn_delay_ms":0,"purpose":"post_workload_readmission","server_commit":"8829189992d9290f1aa128fb00950dd854f43211","client_commit":"2b7fcbf52782f8729fd6ce6c9ab803617d72897b","tp_size":"2","server_max_batch":"2","duration":0.785514358,"conversations_completed":4,"conversations_failed":0,"avg_turns_completed":2.0,"avg_conversation_duration_ms":680.88620825,"completed":8,"failed":0,"total_input_tokens":416,"total_output_tokens":64,"request_throughput":10.184409640033595,"input_throughput":529.5893012817469,"output_throughput":81.47527712026876,"total_token_throughput":611.0645784020156,"mean_ttft_ms":189.25433775,"median_ttft_ms":236.0510255,"std_ttft_ms":91.00097739492537,"p50_ttft_ms":236.0510255,"p90_ttft_ms":262.8936926,"p99_ttft_ms":264.10024466,"mean_tpot_ms":21.52321925,"median_tpot_ms":21.59337264285714,"std_tpot_ms":0.9822315518907105,"p50_tpot_ms":21.59337264285714,"p90_tpot_ms":22.48637407142857,"p99_tpot_ms":22.572006821428573,"mean_itl_ms":16.739350569444447,"median_itl_ms":20.280244,"std_itl_ms":9.372785774041324,"p50_itl_ms":20.280244,"p90_itl_ms":21.8494206,"p99_itl_ms":30.806120210000003,"mean_e2el_ms":339.9168725,"median_e2el_ms":388.02185349999996,"std_e2el_ms":86.74422537684215,"p50_e2el_ms":388.02185349999996,"p90_e2el_ms":401.94781559999996,"p99_e2el_ms":413.76534426,"per_turn_metrics":[{"turn_index":0,"num_samples":4,"completed":4,"failed":0,"total_input_tokens":160,"total_output_tokens":32,"request_throughput":5.092204820016797,"input_throughput":203.68819280067189,"output_throughput":40.73763856013438,"total_token_throughput":244.42583136080626,"mean_ttft_ms":127.99577774999999,"median_ttft_ms":128.1248625,"std_ttft_ms":94.30425103527772,"p50_ttft_ms":128.1248625,"p90_ttft_ms":226.86030060000002,"p99_ttft_ms":234.47434056,"mean_tpot_ms":22.089172857142856,"median_tpot_ms":22.114122285714284,"std_tpot_ms":0.43515371919893253,"p50_tpot_ms":22.114122285714284,"p90_tpot_ms":22.540744071428573,"p99_tpot_ms":22.57744382142857,"mean_itl_ms":17.17952794444444,"median_itl_ms":20.2898775,"std_itl_ms":9.599872306848535,"p50_itl_ms":20.2898775,"p90_itl_ms":25.1509305,"p99_itl_ms":30.89908985,"mean_e2el_ms":282.61998775,"median_e2el_ms":282.5744265,"std_e2el_ms":91.4721237188371,"p50_e2el_ms":282.5744265,"p90_e2el_ms":378.84381890000003,"p99_e2el_ms":386.90337479},{"turn_index":1,"num_samples":4,"completed":4,"failed":0,"total_input_tokens":256,"total_output_tokens":32,"request_throughput":5.092204820016797,"input_throughput":325.901108481075,"output_throughput":40.73763856013438,"total_token_throughput":366.6387470412094,"mean_ttft_ms":250.51289774999995,"median_ttft_ms":250.51778950000002,"std_ttft_ms":12.80006499169254,"p50_ttft_ms":250.51778950000002,"p90_ttft_ms":263.6597574,"p99_ttft_ms":264.17685114,"mean_tpot_ms":20.957265642857145,"median_tpot_ms":21.52421721428572,"std_tpot_ms":1.0486142822079907,"p50_tpot_ms":21.52421721428572,"p90_tpot_ms":21.611059014285715,"p99_tpot_ms":21.63493561571428,"mean_itl_ms":16.299173194444446,"median_itl_ms":20.267008999999998,"std_itl_ms":9.118835768968623,"p50_itl_ms":20.267008999999998,"p90_itl_ms":21.07201,"p99_itl_ms":29.13940425,"mean_e2el_ms":397.21375724999996,"median_e2el_ms":392.7659,"std_e2el_ms":10.775242489752332,"p50_e2el_ms":392.7659,"p90_e2el_ms":409.4510084,"p99_e2el_ms":414.51566354}]} diff --git a/docs/index.md b/docs/index.md index c3a584417..e874fdcff 100644 --- a/docs/index.md +++ b/docs/index.md @@ -46,7 +46,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | | `models/qwen35/roadmap.md` | Qwen3.5 dense roadmap v2 (#654): core correctness/admission/chunked-prefill/sampling/step-tail gates are landed; current 4B HTTP boundary is the retained #469 RTX 5090 sweep, which completed with zero failed requests but trails vLLM at high concurrency. Next: HTTP gap attribution, mixed-load ITL (#470), lifecycle recovery (#471), joint-state prefix reuse (#257), and design-first TP (#446). | -| `models/qwen35/load-snapshot.md` | Issue #605 publishes Qwen3.5 logical running, waiting, and KV load through the shared single-GPU/TP scheduler backend and `EngineHandle::with_load_watch`. | +| `models/qwen35/load-snapshot.md` | Qwen3.5 publishes logical running, current pending, and KV load after drain/cancellation pruning and before admission through the shared single-GPU/TP scheduler backend. | | `models/qwen35/prefix-cache.md` | Qwen3.5 prefix-cache design: a hit is valid only when full-attention KV and a complete recurrent/conv snapshot exist at the same 256-token boundary; the first version uses a fixed-budget GPU snapshot pool with joint lookup, pinning, and LRU eviction. | | `models/qwen35/kv-admission.md` | Issue #254 complete: Qwen3.5 now uses full-lifetime KV admission, deferred pressure handling, impossible-request rejection, explicit error semantics, direct rejection-event coverage, RTX 5090 e2e, and real HTTP pressure/post-pressure validation. | | `models/qwen35/optimization.md` | Hybrid 24 linear + 8 full attn optimization ledger. Decode-tuning refresh fuses MLP gate/up and tunes decode cublasLt buckets, improving direct TPOT by 2-3%; vLLM still leads 1024/256 HTTP decode. | @@ -54,7 +54,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (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/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/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 and P2A are complete: TP2 has start-gated eager unified prefill+decode, strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load ordinal validation; P2B GDR state sharding is next. | | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. | | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. | | `models/qwen35/unified-prefill-overlap.md` | Issue #715 implementation record: opt-in single-GPU shared-SM overlap keeps one prefill chunk in flight while active decode continues; default serial policy and unsupported-combination guards remain explicit. | @@ -224,6 +224,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `benchmarks/qwen3-4b-serving-vllm-rtx5090.md` | **Deleted** — superseded by `models/qwen3/serving-performance.md`. | | `benchmarks/k3-ep4-decode-profile.md` | K3 EP4 93-layer decode step profile (MegaMoE, GB300, 2026-08): 52% backbone B=1 dense GEMM at ~50% of SOL (cuBLASLt splitK), 20% fused MoE, 5% a mis-tuned router top-k; not launch-bound. Levers: near-SOL dense GEMV (~10 ms) + router fix (~2 ms) project ~50 → ~30 ms/step; FP8 backbone is the structural next step. | | `benchmarks/deepseek-v2-lite-vllm-tp2-ep2.md` | DeepSeek-V2-Lite EP2 2026-06-28 snapshot: PegaInfer host-staged/NCCL passed correctness, direct diagnostics, HTTP pressure, and trace rows; stock vLLM TP2/TP2+EP2 are retained as FlashInfer SM120/CUDA 12.8 setup failures, with a separate FlashInfer-fixed vLLM validation and no parity claim. | +| `benchmarks/qwen35-tp2-phase2a-multiturn.md` | Qwen3.5 Phase 2A real TP2 multi-turn serving gate: 12/12 dependent conversations and 44/44 turns completed under concurrency above server capacity, followed by a no-restart 4/4-conversation readmission probe; exact commands, pins, per-turn history growth, pass criteria, and raw JSON are retained. | | `benchmarks/qwen35-4b-serving-vllm-rtx5090-2026-07.md` | Qwen3.5-4B vs vLLM 0.25.1 on 1x RTX 5090 for #469: correctness gates and retained HTTP matrix completed with zero failed requests, but PegaInfer does not reach vLLM parity; requested 1024/256 c16 is `17.36ms` / `807 tok/s` vs vLLM `9.34ms` / `1425 tok/s`, while direct c16 TPOT `9.14ms` points first to HTTP/frontend/scheduler attribution. | | `benchmarks/qwen35-4b-serving-vllm-rtx5090.md` | Qwen3.5-4B TP1 vs vLLM 0.23.0 on RTX 5090: latest direct PegaInfer A/B improves TPOT by 2-3%; HTTP `vllm bench serve` shows prompt-len-1 decode close, but vLLM still leads 1024/256 TPOT and high-concurrency output tok/s. Includes Nsight Systems direct/HTTP gap notes. | | `benchmarks/qwen-mixed-sampling-http.md` | Issue #412 HTTP mixed-sampling evidence: Qwen3-4B `/v1/completions` completed 64/64 with 32 greedy + 32 sampled requests, failed=0/timeouts=0, TTFT/TPOT/ITL/output tok/s retained; Qwen3.5-4B passed the same workload as supplemental evidence. | diff --git a/docs/models/qwen35/load-snapshot.md b/docs/models/qwen35/load-snapshot.md index 876196c0a..dcbad0ab7 100644 --- a/docs/models/qwen35/load-snapshot.md +++ b/docs/models/qwen35/load-snapshot.md @@ -1,8 +1,8 @@ # Qwen3.5 Scheduler LoadSnapshot -> **TL;DR:** Qwen3.5 publishes one logical `LoadSnapshot` stream from its shared single-GPU/TP scheduler: running counts active and prefilling requests, waiting counts deferred requests, and KV usage is request-page capacity minus available pages. +> **TL;DR:** Qwen3.5 publishes one logical post-drain/post-prune `LoadSnapshot` stream from its shared single-GPU/TP scheduler: running counts active and prefilling requests, waiting counts all current pending work, and KV usage is request-page capacity minus available pages. > -> **Last touched:** 2026-07 +> **Last touched:** 2026-08 ## Preparation @@ -35,28 +35,29 @@ Qwen3.5 SchedulerBackend Both Qwen3.5 execution modes own one logical request stream, so single-GPU and TP each attach one `EngineHandle::with_load_watch` receiver. The frontend bridge, metric names, labels, and scheduler-stat conversion remain unchanged. -The scheduler publishes at the top of its existing loop. At that point, work retired by the previous step has been removed and its KV pages have been released, so the next snapshot can settle to idle before `blocking_recv()` waits for new work. +Each scheduler tick first merges deferred work with every submission currently available, then prunes closed pending, active, and prefilling requests before publishing. The fixed boundary is `drain -> prune -> publish load -> admission -> plan`. If the idle scheduler wakes through `blocking_recv()`, it drains, prunes, and publishes again before admission so work closed before admission never consumes a slot or appears in the snapshot. Snapshot accounting is: | Metric field | Existing Qwen3.5 state | | --- | --- | | `num_running_reqs` | `active.len() + prefilling.len()` | -| `num_waiting_reqs` | `deferred.len()` | +| `num_waiting_reqs` | the merged pending queue: prior deferred work plus newly drained submissions | | `kv_used_blocks` | request KV capacity minus currently available request pages | | `kv_total_blocks` | backend request KV capacity, excluding the CUDA Graph padding page | -Publication reads the scheduler's existing queues and KV allocator at the settled boundary. +Publication reads the scheduler's queues and KV allocator after closed resident state has gone through its normal retirement path. The snapshot therefore describes the state used by the following admission decision: cancelled residents no longer count as running or hold capacity, while live pending requests count as waiting even if they were submitted during the current tick. The live gate uses `scripts/bench_http_serving.py` to create overlapping HTTP traffic and a 100 ms `curl /metrics` sampler to retain the three labeled gauges. ## Execution Log - Added load watches to `start_with_capacity` and `start_tp_with_capacity` and attached each receiver to its engine handle. -- Added direct, backend-neutral `LoadSnapshot` publication at the top of the shared scheduler loop, following Qwen3's instrumentation shape. +- Added direct, backend-neutral `LoadSnapshot` publication in the shared scheduler loop, following Qwen3's instrumentation shape. - Derived KV capacity and availability through `SchedulerBackend`, so the same publication logic serves single-GPU and TP. - Validated the single-GPU path with the existing scheduler E2E and live HTTP pressure: running and KV usage rose during generation, waiting reached three at `--max-batch 1`, and every gauge returned to zero after drain and recovery. - Updated the shared Prometheus documentation for Qwen3.5's one-logical-engine contract. +- P2A step 2 moved publication after drain and cancellation pruning. Focused CPU tests cover closed pending and resident work, and a real TP1 `max_batch=1` gate proves a cancelled resident disappears from the post-prune load, frees capacity for same-tick admission, and leaves running, waiting, and KV usage at zero after recovery. ## Validation Boundary @@ -155,10 +156,10 @@ The server exited cleanly and the metric sampler reported no errors. ## Debrief -- **Outcome**: Qwen3.5 now feeds one logical `LoadSnapshot` stream to the frontend for both single-GPU and TP. Single-GPU live validation confirms running, waiting, KV usage, idle reset, and recovery; TP uses the same scheduler path but was not part of this live run. +- **Outcome**: Qwen3.5 feeds one logical `LoadSnapshot` stream to the frontend for both single-GPU and TP. Publication now occurs after current submissions are drained and closed work is pruned, so the snapshot is the admission boundary rather than a view of only the previous tick. - **Pitfalls encountered**: - The TP scheduler rebase required KV accounting through `SchedulerBackend`; retaining model-specific `model.kv_pool()` access would not compile against the shared loop. - **Lessons learned**: - A shared scheduler loop should expose observability through `SchedulerBackend` so one implementation covers both execution topologies. - - Top-of-loop publication captures settled state, including the idle zero after KV pages return. + - Post-drain/post-prune publication captures both newly waiting work and capacity returned by cancellation before the next admission decision. - The existing HTTP benchmark plus raw metric sampling covers the live gauge contract. diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index b2519f936..27d90b170 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,8 +1,8 @@ # 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. +> **TL;DR:** Qwen3.5 TP Phase 2 is two separately delivered correctness milestones: P2a adds eager `RunUnifiedStep` with a shared ordered `RequestId` plan while retaining Phase 1 replicated GDR; P2b shards the head-indexed linear-attention/GDR surface and adds only the hidden all-reduce after local `out_proj`. > -> **Last touched:** 2026-06 +> **Last touched:** 2026-08 ## Goal @@ -18,7 +18,7 @@ Reuse the Qwen3 TP shape: - `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 +- rank-local CUDA context, cuBLAS, and NCCL resources - hidden all-reduce after row-parallel projections - replicated embedding/lm_head as the first-pass simplification @@ -44,20 +44,22 @@ These decisions are settled before implementation starts. - 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. +- `DropRequest`, finish cleanup, cancellation cleanup, and client disconnect must release or reset the corresponding rank-local KV/recurrent/conv state on every rank by `RequestId`. - 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 +## Phase 2 Delivery Boundaries -These topics should not block Phase 1 eager dense TP, but they remain design work before any later implementation. +P2a and P2b are separate implementation series. P2a must complete its protocol and lifecycle gates before P2b changes loader, kernel, or state shapes. This keeps worker-protocol failures distinguishable from local-head loader/kernel/state failures. -- 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`. +The following are separate follow-up RFCs, not Phase 2 deliverables: + +- TP CUDA Graph capture/replay, including graph slots, padding, synchronized capture, and recurrent/conv D2D compaction. +- TP-aware prefix caching and recurrent-state snapshots. +- Vocabulary-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. + +Phase 2 is a correctness and per-rank HBM-reduction milestone. It makes no speedup promise: P2b adds one hidden all-reduce in each of the 24 linear-attention layers. Report a matched Phase 1 TP2 versus P2b TP2 A/B before making any performance claim. ## Why Dense First, GDR Second @@ -65,7 +67,7 @@ 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. +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 `DropRequest` cleanup plus re-admission must preserve rank-local recurrent-state boundaries. 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. CUDA Graph slot, padding, and compaction semantics remain a separate follow-up RFC. ## Architecture Summary @@ -166,7 +168,7 @@ State ownership: - 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 +- all workers must observe the same ordered `RunPrefillChunks`, `RunDecodeStep`, `DropRequest`, and `Shutdown` commands CUDA Graph: @@ -180,35 +182,71 @@ Validation scope: - 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 +- finish, explicit drop, cancellation, and client-disconnect cleanup by `RequestId` +- subsequent admission with a new `RequestId` observes no stale KV, 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 +## P2a: Eager TP Unified Execution -Phase 2 converts linear attention from replicated execution to true TP execution. +P2a implements eager TP `RunUnifiedStep` while retaining the Phase 1 replicated linear-attention/GDR weights, kernels, conv state, recurrent state, and scratch shapes. -Shard: +Every rank receives the same canonical `UnifiedPlan`: ordered prefill and decode items, each carrying a `RequestId` and the request-local execution inputs needed for that row. The order is the collective-order contract; every rank executes the actual plan rows in that order. P2a does not introduce CUDA Graph padded-slot semantics, D2D state movement, or a cross-rank slot-compaction protocol. -- `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` +State and artifacts are keyed by `RequestId`: -Execution: +- worker-local KV, conv, and recurrent state is found, created, promoted, and released by `RequestId`; +- the primary worker returns prefill and decode artifacts carrying their `RequestId`; +- the scheduler resolves artifacts by ID and rejects unknown, duplicate, or missing results instead of relying on returned row position; +- finish, explicit drop, cancellation, and client disconnect broadcast the same `DropRequest(RequestId)` lifecycle command to every rank; +- each worker returns `DropAck { existed: bool }`, where `existed` reports whether that rank actually found and removed the request state; +- every controller-side drop carries a `DropExpectation::{MustBeAbsent, MustExist}` derived from scheduler-owned lifecycle state, not from a separate `worker_state_materialized` flag; +- cancellation before the first successful prefill dispatch requires `MustBeAbsent` and an exact-rank all-false `DropAck` set; partially prefetched, active, disconnected, and completion-candidate requests require `MustExist` and an exact-rank all-true set; +- uniformity alone is insufficient. All-false for `MustExist`, all-true for `MustBeAbsent`, or mixed true/false values prove lifecycle divergence and poison the whole TP executor even if the drop leaves every rank absent. Dispatch failure, a missing acknowledgement, or a rank-local drop failure is likewise fatal; the scheduler must not warn and continue promotion or serving. + +The scheduler allocates a TP `RequestId` before the first prefill command materializes worker state, so cancellation in that interval legitimately requires `MustBeAbsent`; every successfully prefetched or active request requires `MustExist`. If execution mutates state and then fails, the executor is already poisoned and no healthy-path drop is attempted. Worker replies prove cross-rank existence, while scheduler lifecycle determines which value is valid. + +Successful request completion has a fail-closed commit boundary. After computing a terminal artifact, the scheduler keeps the request unresolved in a local completion candidate and withholds its user-visible terminal events. An EOS candidate buffers only `Finished`; a length-limited candidate buffers the final `Token` followed by `Finished`; completion on the first prefill token, including `max_tokens <= 1`, follows the same rule. A completion candidate is materialized, so the scheduler first requires a valid `MustExist` all-rank all-true `DropAck` contract, then removes the logical request from scheduler state, and only then publishes the buffered events in order. A client that observes `Finished` can therefore rely on consistent rank-local cleanup acknowledgement for that request. + +If completion drop fails, the scheduler publishes neither the buffered final `Token` nor `Finished`, poisons the replica, and keeps the candidate unresolved for the complete terminal error fan-out. The client may already have received earlier streamed tokens, so this boundary makes only successful termination atomic; it does not make the complete streamed response transactional. The completion candidate is a scheduler-local prepare/commit abstraction, not a distributed `ValidateUnified -> ExecuteUnified` protocol. A process-level fail-stop may still prevent both the acknowledgement and the terminal `Error` from being delivered. + +Controller response collection validates an exact rank set, not only a message count. Every response rank must be in `0..world_size` and appear exactly once. Ping requires `Ack` from every rank; drop requires a `DropAck` from every rank whose existence values match the controller's `DropExpectation`; prefill, decode, and unified execution require exactly one matching typed result from rank 0 and `Ack` from every non-primary rank. Missing, duplicate, out-of-range, wrong-variant, mixed drop-existence, or uniformly unexpected drop-existence responses are protocol failures. If workers may already have mutated state, any such response-set failure poisons the complete TP replica. + +Cancellation cleanup is a scheduler-tick boundary, not a late planning cleanup. At the start of each tick the scheduler first merges deferred work with newly received submissions, then prunes every already-closed sink from active, prefilling, and pending work. Healthy active/prefilling TP removals complete their all-rank `DropRequest` before the scheduler publishes load or computes admission capacity; pending requests have no worker state and are discarded directly. The required order is `drain -> prune -> publish load -> admission -> plan`. This keeps cancelled requests out of running/waiting metrics, decode slots, future-KV reservation, and the current tick's prefill budget. A cancellation racing after the prune check is still retired by the existing token-send failure path. + +Artifact alignment follows the ordered plan rather than returned row position. Prefill expects artifacts only for `finish_prefill == true` IDs and preserves explicit absence for non-final rows; decode expects one artifact for every row. Unknown, duplicate, or missing IDs are fatal after execution. This TP adapter contract does not require the TP1/single-GPU logits-and-sampling path to adopt sparse artifacts. + +P2a separates recoverable plan rejection from fatal replica failure. Controller-provable structural errors reject before dispatch without poisoning. Worker-local existence, phase, materialization, or capacity mismatches are not protected by an all-rank validation barrier and are therefore replica-fatal, as is any CUDA/NCCL, response-set, artifact, or lifecycle failure after execution is released. + +When fatal failure returns control, the scheduler closes and drains submissions, emits exactly one terminal `Error` for every unresolved stable or tick-local request, publishes an idle load snapshot, exits, and begins whole-executor teardown. Exactly-once fan-out comes from consuming mutually exclusive request owners; pre-admission requests do not yet have a TP `RequestId`, so that ID is not a global deduplication key. The scheduler does not retry per-request drops or claim rank-local cleanup after poison. Collective or teardown timeout remains process-level fail-stop, where client fan-out cannot be guaranteed; rollback, communicator recovery, and in-place restart remain out of scope. + +The worker's internal request table may remove and reinsert entries as an implementation detail, but that is not CUDA Graph slot semantics and must not require state copying between request identities. + +The fixed 16-device Triton AOT handle table remains. Before model loading or worker launch, TP startup validates every requested logical CUDA ordinal against that supported range; dynamic handle allocation is not a P2a prerequisite. `tp_size > 1` with CUDA Graph requested continues to fail closed. + +P2a acceptance retains the Phase 1 TP1/TP2 gates and adds TP2 mixed chunked-prefill/decode, cleanup and re-admission, strict artifact-ID checks, and fatal-path coverage. Drop tests must accept only `MustBeAbsent`/all-false and `MustExist`/all-true exact-rank sets, rejecting mixed or uniformly unexpected existence and malformed rank/reply sets. Completion tests must withhold EOS, length, and immediate-prefill success events until `MustExist` cleanup succeeds; controller-rejected plans must preserve executor health, while worker-local or post-execution failures must stop the scheduler and fan out terminal errors to every accepted unresolved request. All work remains eager and uses the replicated GDR path. + +Lifecycle observability, cancellation ordering, fail-closed cleanup, and unified execution remain separately reviewable implementation boundaries. Cancellation ordering fixes a pre-existing TP1 scheduler issue but stays in P2a because unified planning depends on the pruned state. + +## P2b: Local-Head Linear Attention / GDR + +P2b converts the 24 linear-attention layers from replicated execution to true TP execution. It additionally requires `linear_num_key_heads % tp == 0` and `linear_num_value_heads % tp == 0`; unsupported degrees and unsupported local kernel shapes fail before model loading. + +Shard every head-indexed linear-attention/GDR surface by the local key/value-head ranges: + +- `in_proj_qkv`, preserving local q/k/value channel layout; +- `in_proj_z`, `in_proj_b`, `in_proj_a`, `conv1d_weight`, `dt_bias`, and `A_log`; +- `out_proj` input columns; +- conv state, recurrent state, GDR scratch, and intermediate buffers. + +The head-dimension `norm_weight` remains deliberately replicated because it is shared by every local value head; it is not a head-indexed state or collective surface. Embedding and tied `lm_head` also remain replicated. -- 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` +Each rank runs local projections, convolution, GDR prefill/decode kernels, gated RMSNorm/output-gate work, and local `out_proj` against local dimensions. The only linear-attention collective is the hidden all-reduce after `out_proj`. Conv state and GDR recurrent state are request-local and rank-local for their full lifetime and are never all-reduced or centralized. -Never all-reduce GDR recurrent state or conv state. Their ownership is rank-local and request-local. +P2b acceptance requires loader reconstruction/layout tests, local AOT-kernel shape validation, rank-local allocation checks, Phase 1 and P2a regression gates, short and long TP2 HF replay, cleanup without stale request state, and a matched Phase 1 TP2 versus P2b TP2 HBM/latency/throughput report. The report is evidence, not a speedup threshold. ### vLLM Reference @@ -221,21 +259,21 @@ Use vLLM's `Qwen3NextForCausalLM` / `QwenGatedDeltaNetAttention` as the referenc - 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 -PegaInfer-specific work remains: worker-owned rank-local recurrent state, `RequestId` lifecycle, local-state slot compaction, `DropRequest` cleanup, and fail-closed kernel-shape validation. +PegaInfer-specific work remains: worker-owned rank-local recurrent state, `RequestId` lifecycle, request-state removal and re-admission, `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 +- request-state cleanup and re-admission replay +- recurrent-state cleanup on finish/drop/cancellation +- no stale local recurrent state after a new `RequestId` is admitted ## References - `docs/models/qwen3/tp-design.md` -- `pegainfer-qwen3-4b/src/config.rs` -- `pegainfer-qwen3-4b/src/executor.rs` +- `pegainfer-qwen3/src/config.rs` +- `pegainfer-qwen3/src/executor.rs` - `pegainfer-qwen35/src/config.rs` - `pegainfer-qwen35/src/weights.rs` - `pegainfer-qwen35/src/recurrent_state.rs` diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 66d35cd26..d8ebda80d 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -1,8 +1,8 @@ # 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. +> **TL;DR:** Qwen3.5 TP Phase 1 and P2A are complete: TP2 now supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation; P2B GDR state sharding is next. > -> **Last touched:** 2026-07 +> **Last touched:** 2026-08 ## Scope @@ -15,6 +15,7 @@ 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 +- prompt echo support; Qwen3.5 accepts completion-only requests ## Phase 1 Outcome @@ -136,21 +137,305 @@ Stable test knobs: - `PEGAINFER_TEST_TP_DEVICES`: comma-separated TP2 CUDA ordinals. Defaults to `0,1`; examples: `1,2`, `2,3`. TP2 tests require exactly two distinct ordinals. - `PEGAINFER_TEST_FRONTEND_MODEL_PATH`: optional tokenizer/config metadata path for HTTP serving tests. Defaults to `PEGAINFER_TEST_MODEL_PATH` when unset. -## Follow-Up Work +## Phase 2 Progress -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. +Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, so P2b can now change loader, kernel, and state shapes while preserving those contracts. -### TP mixed-step unified execution +### P2a: 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. +P2A implements eager `RunUnifiedStep` under TP while retaining Phase 1's replicated linear-attention/GDR weights, kernels, conv state, recurrent state, and scratch shapes. -Goals: +#### Step 1: lifecycle cleanup gates + +Completed as test and fault-observation infrastructure before changing production lifecycle semantics. + +Why it exists: + +- A successful controller-side `DropRequest` return did not directly prove that every rank removed the same `RequestId` or released its KV/recurrent/conv ownership. +- Later cancellation, partial-dispatch, drop-acknowledgement, and unified-step work needs direct evidence of rank-local state before and after each transition. +- Cleanup must prove both capacity recovery and fresh numeric state; scheduler bookkeeping alone cannot detect a stale rank-local recurrent or KV allocation. + +Implemented in `tp_executor.rs` under `#[cfg(test)]` only: + +- `WorkerStateSnapshot { rank, request_count, requests }`, where each request entry carries its `RequestId` and `Prefilling`/`Decoding` phase. +- A healthy snapshot API that requires an unpoisoned executor and an unchecked test-only API that bypasses only the controller poison guard. The latter is reserved for later synthetic failure tests while every worker channel remains connected. +- Exact-rank collection that rejects duplicate or out-of-range ranks, payload/response rank disagreement, missing responses, wrong reply variants, and inconsistent request counts. Valid snapshots are returned in rank order. +- An ignored TP2 capacity gate that fills configured `max_batch=2`, observes every ID in `Decoding` on both ranks, drops every ID, requires zero requests on every rank, refills the complete capacity, executes decode, and proves the second cleanup is also empty. +- An ignored TP2 numeric gate that records a prompt's deterministic first token and five requested logprobs, drops the request, verifies every rank is empty, re-admits the same prompt under a new `RequestId`, and requires an exact artifact match. + +Verification: + +- Qwen3.5 release all-target check and clippy with `-D warnings` pass. +- The regular library suite passes: `71 passed`, `0 failed`, `8 ignored`. +- Both new TP2 GPU gates pass independently. +- Formatting and `git diff --check` pass. + +This step deliberately adds no production command, scheduler state transition, drop semantics, or serving claim. It supplies the observability needed to prove the later P2A changes rather than trusting controller-visible acknowledgements alone. + +#### Step 2: cancellation pruning before admission + +Completed as the first production scheduler lifecycle change in P2A. + +The shared TP1/TP scheduler tick now follows the fixed order `drain -> prune -> publish load -> admission -> plan`. It merges deferred and newly submitted work, removes requests whose `TokenSink` is already closed, publishes the resulting state, and only then computes slot/KV budgets and admits work. The idle wakeup path repeats drain/prune/publication after `blocking_recv()` because the waking request can close before admission and other submissions can race with the receive. + +Change to the existing scheduler skeleton: + +- Previously, the loop published the state left by the prior tick before draining `submit_rx`; `num_waiting_reqs` therefore represented only `deferred.len()`. It then drained submissions and entered admission. +- The loop now first takes `deferred`, drains every currently available submission into the same `pending` vector, and calls one backend-neutral `prune_closed_requests` helper across pending, active, and prefilling ownership. +- `publish_load` remains the same internal watch publication surface and `LoadSnapshot` retains the same public fields and types. A small `logical_load_counts` helper makes the post-prune running/waiting calculation directly testable; waiting is now `pending.len()` at the admission boundary. +- The empty-loop blocking skeleton is retained, but wakeup now has a second drain/prune/publish boundary before admission. If every received request is already closed, the loop returns to blocking without entering admission or planning. +- The admission and plan implementations are otherwise unchanged. This commit changes which settled state they consume, not their policies or the public `EngineHandle::load_watch()` API. + +Effect on the pre-existing TP1 (`SchedulerBackend::Single`) path: + +- A closed pending request is removed before `alloc_prefill_state`, so TP1 no longer allocates its KV/recurrent state and then discovers the disconnected consumer during token delivery. +- A closed prefilling request is removed before plan construction. Consuming its existing `PrefillBackendState::Single` drops the owned `KvState` and `RecurrentState` through their existing RAII path; no TP1 release routine or allocator rule was added. +- A closed active request is retired before the next model step through the existing `compact_single_slot` path. The same `swap_remove` and graph-slot copy used by EOS, length, and token-send failure are reused; only the retirement timing moves earlier, so the cancelled row no longer executes and samples one unnecessary decode step. +- TP1 page availability, slot budgets, and plan construction now observe that cleanup in the same tick. Consequently, the executed batch width and load samples can differ under cancellation, but KV accounting formulas, slot-compaction mechanics, admission policy, plan builders, kernels, and sampling implementation are unchanged. + +Cleanup deliberately reuses existing backend ownership paths: + +- pending cancellation uses stable `retain`, preserving FIFO order among live requests; +- active cancellation uses normal request retirement, including TP drop and single-GPU slot compaction; +- prefilling cancellation removes the entry and drops its backend state through the existing adapter. + +This ordering makes cancellation visible to admission in the same tick. A closed resident is absent from the post-prune running/KV state, a live replacement is present in waiting, and the capacity released by the resident can admit that replacement immediately. + +Verification: + +- Three focused CPU tests pass for pending FIFO pruning, post-prune logical load, and same-tick capacity reuse. +- The ignored real TP1 `max_batch=1` cancellation/replacement gate passes on an RTX 3090. It observes the post-prune `running=0, waiting=1` boundary, completes the replacement, and returns running, waiting, and KV usage to zero. +- The existing real TP1 scheduler integration E2E passes (`1 passed`, 29.06 seconds). +- Release all-target clippy with `-D warnings` passes. +- The regular library suite passes: `74 passed`, `0 failed`, `9 ignored`. + +This step does not strengthen the TP drop reply protocol or add scheduler-wide fail-closed propagation. It uses the Phase 1 cleanup calls as they exist; exact-rank `DropExpectation` acknowledgement belongs to step 3, and mandatory replica-fatal propagation belongs to step 4. + +#### Step 3: TP worker protocol hardening + +Completed as the distributed command/reply foundation for later scheduler fail-closed handling and unified execution. + +State-mutating prefill, decode, drop, and future unified envelopes now share one `Pending -> Execute | Cancel` start gate. The controller enqueues an envelope for every rank before resolving `Execute`; if any enqueue fails, it resolves `Cancel` for the delivered prefix and poisons the executor. Workers wait on the gate before request-state mutation, kernel launch, or NCCL entry. Ping and test-only snapshots remain ungated because they do not mutate request state or enter collectives. + +Response handling now separates timed transport collection from pure response-set validation. Ping, prefill, decode, and drop all require exactly one in-range response from every rank: + +- ping requires `Ack` from every rank; +- prefill requires one `Prefill` result from rank 0 and `Ack` from every non-primary rank; +- decode requires one `Decode` result from rank 0 and `Ack` from every non-primary rank; +- drop requires `DropAck { existed }` from every rank and validates the complete existence vector against the scheduler's `DropExpectation`. + +Duplicate, missing, out-of-range, wrong-variant, non-primary typed, or missing-primary responses poison the executor after dispatch. Drop accepts only exact-rank all-false for `MustBeAbsent` and exact-rank all-true for `MustExist`; mixed values and uniformly unexpected values are lifecycle divergence even though every rank is absent after the command. + +The low-level `Qwen35TpExecutor::drop_request` API now requires `DropExpectation` and returns `Result<()>`; `DropExpectation` is re-exported through the model-local `runtime` module for tests and debugging. The server-facing `EngineHandle` API is unchanged. Scheduler callers derive expectations from owned lifecycle state: + +- active and successfully dispatched/final-prefill state use `MustExist`; +- prefilling cancellation at `cursor == 0` uses `MustBeAbsent`; +- prefilling cancellation after progress uses `MustExist`; +- execution failure performs no healthy-path drop because the executor is already poisoned. + +This commit intentionally retains the temporary scheduler adapter that logs a returned healthy-drop failure. Executor poison prevents subsequent normal commands, but step 4 still owns mandatory propagation into one scheduler-wide terminal path, fail-closed completion publication, and exactly-once unresolved-request error fan-out. + +Verification: + +- Release all-target check and clippy with `-D warnings` pass. +- Pure protocol/gate tests cover exact-rank reply matrices, lifecycle-expectation mismatch, controller structural rejection with zero dispatch, all-enqueued `Execute`, prefix-only `Cancel`, and poison preservation. +- The complete regular library suite passes on SM86: `82 passed`, `0 failed`, `12 ignored`. +- Three new real TP2 gates pass: prefix-only dispatch failure leaves every rank empty (214.38 seconds under unrelated four-GPU training contention), lifecycle expectations plus mixed-rank divergence pass (8.05 seconds), and a real receiver disconnect poisons without claiming an unavailable all-rank snapshot (6.96 seconds). +- Existing healthy TP2 prefill/decode/drop and scheduler chunked-prefill/decode smoke tests pass (6.45 and 6.54 seconds). + +Implementation pitfall: the first release-only gate test hung because `start.execute()` was placed inside `debug_assert!`; release builds remove the complete assertion expression, including side effects. The state transition now executes unconditionally and only its boolean result is debug-asserted. Protocol state transitions must never live inside debug-only assertions. + +#### Step 4: fail-closed TP scheduler lifecycle + +Completed as the scheduler recovery and user-visible completion boundary required before TP unified execution. + +The temporary step-3 logging adapter is gone. TP active retirement and prefill-state cleanup now return `Result`, and cancellation pruning, standalone prefill/decode execution, token dispatch, and final-prefill promotion propagate any returned TP lifecycle failure to the scheduler loop. A returned TP prefill/decode or post-execution artifact-alignment error follows the same fatal path. The scheduler does not retry per-request cleanup after poison. + +Successful TP completion is now prepared and committed explicitly: + +- decode EOS buffers only `Finished(Stop)`; +- decode length buffers the final `Token` followed by `Finished(Length)`; +- immediate final-prefill EOS/length uses the same buffering rule, so its first token is not exposed before cleanup; +- the candidate keeps logical request ownership while `DropRequest(MustExist)` runs, and publishes the buffered events only after exact-rank all-true `DropAck` succeeds; +- drop failure discards the buffered success events and transfers the unresolved candidate to terminal error fan-out. + +Fatal errors carry any tick-local request ownership back to the scheduler loop instead of consuming it in a step helper. The single terminal helper closes `submit_rx`, drains every request whose send completed before close, consumes transient candidates/scheduled work plus active, prefilling, current pending, and deferred owners exactly once, attempts one `TokenEvent::Error` per request, publishes a zero-running/zero-waiting/zero-KV load snapshot, and exits. It does not deduplicate by TP or external request ID; exclusivity comes from moving each request out of its prior owner. + +The existing single-GPU boundary remains explicit. TP completion requires cleanup acknowledgement before publication, while the non-TP decode path retains its original `Token`/`Finished` publication before slot retirement. Single-GPU prefill state still drops through its existing RAII ownership, and `EngineHandle` plus `TokenEvent` public types are unchanged. + +Verification: + +- Focused scheduler lifecycle coverage passes: `16 passed`, `0 failed`, `2 ignored`. It covers TP EOS/length and immediate-prefill event buffering, active/prefill drop failure with no leaked success event, remaining scheduled ownership preservation, prune failure propagation, non-TP publication order, and barrier-controlled submit close/drain fan-out across every owner class. +- Release all-target check and clippy with `-D warnings` pass; formatting and `git diff --check` pass. +- The complete regular library suite passes on SM86 outside the GPU-isolated sandbox: `90 passed`, `0 failed`, `12 ignored`. +- The real-weight TP2 scheduler chunked-prefill/decode smoke and the complete TP2 scheduler E2E pass on physical GPUs 1/2. The complete E2E covers context rejection, greedy/logprobs, sequential/repeated/concurrent requests, consumer drop, and post-drop health. +- A same-session TP1 scheduler E2E attempt could not pass model startup because an unrelated training process left insufficient memory for the loader's default 64-slot graph allocation. No training process was stopped. The non-TP ordering CPU gate and complete regular suite are green; the previously established real TP1 E2E evidence remains unchanged. + +Step 4 does not add unified plans, change worker command payloads, alter TP1 sampling, or modify GDR weights/state/kernel shapes. Step 5 owns `RunUnifiedStep`, strict ID-aligned unified artifacts, ordinal validation, and the final TP1/TP2 regression ladder. + +#### Step 5: TP unified execution + +Completed as the final P2A execution-protocol milestone. When active decode and scheduled prefill coexist, TP now uses the same `plan::build_next_plan` decision as the single-GPU scheduler and emits one start-gated `RunUnifiedStep` to every rank. The canonical plan carries ordered prefill/decode items plus separate seeds. The scheduler selects the decode seed first and the prefill seed second to preserve its prior RNG order; workers execute prefill first and decode second on every rank. Rank 0 returns `TpUnifiedResult`, while every non-primary rank must return `Ack`. + +Internal protocol/API changes: + +- Added crate-private `TpUnifiedPlan`, `TpUnifiedResult`, the complete `RunUnifiedStep` payload, and `TpWorkerReply::Unified`. These are not re-exported through `runtime`; the public server `EngineHandle`, `TokenEvent`, and request contract are unchanged. +- Refactored worker prefill/decode into typed inner row operations so standalone and unified commands share the same eager kernels and sampling behavior. Existing public low-level `execute_prefill`/`execute_decode` results retain their shapes. +- Added controller structural validation before enqueue: both halves must be non-empty, row counts must fit scheduler capacity, prefill/decode IDs must be internally unique and mutually disjoint, and prefill chunks must be non-empty. Worker-local request existence, phase, and actual capacity are revalidated after gate release and remain replica-fatal rather than adding a two-phase validation protocol. +- Extended exact-rank reply validation so unified accepts exactly one rank-0 `Unified` result plus non-primary acknowledgements. Returned worker or reply-set failure continues through the step-4 poison and terminal scheduler path. + +Change to the existing scheduler skeleton: + +- Removed the TP-only `build_eager_only_plan` branch. Both backends now use the normal planner, so active decode plus scheduled prefill becomes `ExecutionPlan::Unified` instead of two serialized scheduler ticks. +- Added a TP-only execution adapter that builds the canonical plan, aligns returned artifacts by `RequestId`, processes decode results first, and only then promotes or requeues prefill. Decode completion can therefore release capacity through `DropRequest(MustExist)` before final-prefill promotion. +- Replaced TP prefill token-`0` placeholders and positional decode matching with explicit artifacts. Prefill alignment is a chunk-length `Vec>`: outer `None` means non-final, while `Some { logprob: None, .. }` is a valid final artifact. Decode requires one artifact for every active ID. Shuffled valid results are accepted; unknown, duplicate, non-final, or missing IDs poison the replica after execution. +- The shared `promote_or_requeue` skeleton now accepts a backend-specific artifact wrapper. The single-GPU branch still consumes its original dense sampled tokens/logprobs, uses the same logits path and RNG calls, and retains its step-4 publish/retire semantics. + +Startup now validates TP CUDA ordinals against the generated Triton AOT handle-table contract before model/filesystem/CUDA access. Ordinals must be distinct and below the named table length `16`; the generated wrapper's runtime bounds check remains defense in depth. + +Verification: + +- Release all-target check and clippy with `-D warnings`, formatting, and `git diff --check` pass. +- Complete non-ignored library suite passes on SM86: `96 passed`, `0 failed`, `14 ignored`. +- Pure tests cover ordinal bounds/duplicates, unified structural rejection plus subsequent healthy dispatch, worker-local missing/phase/capacity failures, exact-rank unified replies, shuffled artifact alignment, unknown/duplicate/missing IDs, and final-without-logprobs versus non-final absence. +- The real-weight TP2 executor mixed-step gate passes (`1 passed`, 7.27 seconds), confirms both ranks retain two decoding requests after the combined operation, and returns both ranks to empty after lifecycle cleanup. +- The deterministic real-weight TP2 scheduler case with `max_batch=2` and `max_prefill_tokens=1` passes (`1 passed`, 7.15 seconds). +- All `13` TP2/lifecycle ignored library gates pass together. The fourteenth ignored test is TP1-only and could not load under the unrelated training process: the loader reported `8101 MB` free versus `3248 MB` prefill scratch plus `6288 MB` recurrent state and minimal KV needs. No training process was stopped. +- Complete TP2 scheduler E2E passes (`1 passed`, 35.89 seconds). +- TP2 short/long HF golden gates pass (`2 passed`): short sequential mean `0.0260`, p99 `0.1081`; short batched mean `0.0267`, p99 `0.1167`; long sequential mean `0.0228`, p99 `0.0689`. +- TP2 OpenAI-compatible HTTP serving smoke passes (`1 passed`, 10.41 seconds), including streaming, non-streaming, concurrent completions, logprobs, and TP+CUDA Graph rejection. + +P2A does not shard GDR weights/state/scratch, add a post-GDR all-reduce, enable TP CUDA Graph, or claim a performance improvement. TP1's earlier real-weight evidence remains applicable, but a same-session full TP1 model E2E rerun is still resource-blocked by external GPU occupancy. + +#### Merge-gate test surface + +The pre-merge review reduced P2A coverage to distributed behavior that cannot be +established by restating local helpers. Removed tests enumerated start-gate +booleans, planner choices, CUDA-ordinal bounds, unified-plan shapes, reply +variants, artifact-alignment variants, snapshot-collector variants, and basic +TP2 startup/prefill/chunk/decode smoke already exercised by higher-level gates. + +The retained P2A library surface is seven real TP2 gates: + +- partial dispatch cancels the delivered prefix before rank-local mutation; +- rank-local lifecycle divergence is detected and poisons the replica; +- a disconnected worker receiver fails closed; +- executor and scheduler mixed prefill/decode both complete; +- drop-all restores complete request capacity; +- re-admission under a fresh `RequestId` reproduces the clean first-token artifact. + +Completion cleanup and terminal error fan-out remain covered by the focused +scheduler lifecycle tests. The production-facing regression ladder remains the +TP2 scheduler E2E, short/long HF logits gates, and OpenAI-compatible HTTP +serving lifecycle test. + +On the local 2x RTX 3090 SM86 fixture, the complete explicit ignored run passed +`11/11`: seven library gates, one scheduler E2E, two HF gates, and one HTTP gate. +The regular release library suite passed `94` tests with those seven real TP2 +tests ignored by default. The HF results remained within the established +tolerances: short sequential/batched mean deltas `0.0260`/`0.0267`, and long +sequential mean delta `0.0228`. + +#### Qwen3.5 feature CI gate + +The CUDA workflow now has separate matrix results for Qwen3.5 compile and +Clippy. Both install the repository's CUDA 13.0.2 toolchain, Python 3.10, and +pinned Triton 3.7.1, then build `pegainfer-qwen35` with the actual `qwen35` +feature enabled for all targets on `sm_80`: + +```text +cargo check --release --locked -p pegainfer-qwen35 --features qwen35 --all-targets +cargo clippy --release --locked -p pegainfer-qwen35 --features qwen35 --all-targets -- -D warnings +``` + +This is a compile/lint merge gate, not a substitute for the ignored real-GPU +TP2 tests above. The same target surface passed locally on SM86 with nightly +2026-07-10 and Triton 3.4.0; both commands compiled the Qwen3.5 Triton AOT path. +The GitHub jobs provide the independent Ubuntu, SM80, and pinned Triton 3.7.1 +verification. + +#### Real TP2 multi-turn serving gate + +The Phase 2A production-path gate starts the real OpenAI-compatible Qwen3.5 +server on two RTX 3090 GPUs with TP=2, eager execution, `max_batch=2`, and +`max_prefill_tokens=64`. A pinned upstream Rust `vllm-bench` client runs +dependent `openai-chat` conversations at concurrency 4, so the workload exceeds +resident capacity and forces cleanup followed by fresh admission. + +The primary workload completed 12/12 conversations and 44/44 measured turns +with zero failures. Conversation lengths varied from 2-5 turns, later prompts +carried accumulated history, and every turn returned its configured 24 tokens. +Without restarting the server, a second 4-conversation, 8-turn probe also +completed with zero failures. Graceful shutdown released both TP devices. + +The exact client/server revisions, build and run commands, request counts, +concurrency, prompt/output lengths, per-turn history evidence, pass criteria, +limitations, and raw JSON results are published in +[`docs/benchmarks/qwen35-tp2-phase2a-multiturn.md`](../../benchmarks/qwen35-tp2-phase2a-multiturn.md). + +#### Serving scope: prompt echo is unsupported + +Qwen3.5 serving does not support `echo=true`. Both the legacy and stepped vLLM +bridges submit `echo: false`, so prompt echo has never been part of the current +HTTP serving contract. Direct engine callers can still construct the shared +request type with `echo=true`; the Qwen3.5 scheduler now rejects those requests +immediately after queue drain and cancellation pruning, before load accounting, +capacity admission, backend-state allocation, or TP command dispatch. + +Qwen3.5 no longer emits `TokenEvent::PromptTokens` after prefill. The shared +event variant remains in `pegainfer-frontend` because other model lines still +use it, but it is sealed off from the Qwen3.5 scheduler path. A focused release +test requires a `Rejected` event with the original prompt length and proves the +request is absent from the vector eligible for backend admission. + +Verification on the local SM86 environment passed: + +- focused unsupported-echo release test: `1 passed`; +- complete Qwen3.5 release library suite: `95 passed`, `0 failed`, `7 ignored`; +- all retained real TP2 executor/scheduler gates on two RTX 3090 GPUs: + `7 passed`, `0 failed`; +- release Clippy with `-D warnings` and formatting checks passed. + +#### Model-backed test fixtures are explicit + +Qwen3.5 model-backed tests no longer embed a developer-local absolute weights +path or fall back to a repository-relative model directory. The shared test +fixture resolver requires `PEGAINFER_TEST_MODEL_PATH`, reads and parses its +`config.json`, and confirms the Qwen3.5 model identity before any model or CUDA +initialization. A missing, non-UTF-8, empty, unreadable, malformed, or non-Qwen3.5 +fixture prints one test-specific `SKIP` diagnostic and returns without GPU work. +The optional `PEGAINFER_TEST_FRONTEND_MODEL_PATH` uses the same validation when +set and otherwise reuses the validated engine fixture. + +The resolver source lives under `tests/common` and is also included only for +crate unit-test builds, so the six retained TP executor gates, the mixed-step +scheduler gate, and the integration targets share the same contract without +adding a production API. Static review confirms that no Qwen3.5 test retains a +private model path or implicit `DEFAULT_MODEL_PATH` fallback. + +Verification on the explicit public-compatible SM86 fixture passed: + +- missing fixture, missing `config.json`, and malformed JSON probes emitted the + expected diagnostics without GPU initialization; +- release all-target compile and Clippy with `-D warnings` passed; +- regular release library suite: `95 passed`, `0 failed`, `7 ignored`; +- retained real TP2 library gates: `7 passed`, `0 failed`; +- TP2 scheduler E2E, short/long HF gates, and HTTP serving gate passed; HF + mean/p99 deltas remained `0.0260/0.1081` sequential, `0.0267/0.1167` + batched, and `0.0228/0.0689` long; +- the pinned multi-turn gate rerun completed `12/12` conversations and `44/44` + turns, followed without restart by `4/4` conversations and `8/8` turns, all + with zero failures; graceful shutdown returned both TP GPUs to idle. + +Delivered constraints: - Support mixed prefill+decode scheduler steps under TP. - Preserve deterministic collective ordering across ranks. - Return mixed prefill/decode artifacts from the primary rank. +- Use the `RequestId`-keyed `UnifiedPlan` and artifact contract in `tp-design.md`; P2a does not introduce CUDA Graph padded-slot or compaction semantics. - Validate finish/drop/client-disconnect cleanup under mixed-step execution. - Keep TP CUDA Graph disabled unless a separate graph design is completed. +- Keep the fixed 16-device Triton AOT handle table and add startup-time validation for unsupported logical CUDA ordinals. Why this should be separated from GDR sharding: @@ -158,9 +443,9 @@ Why this should be separated from GDR sharding: - Sharded linear-attention/GDR is a model-state-shape problem. - Combining them would make failures hard to attribute. -### Sharded linear-attention/GDR state +### P2b: sharded linear-attention/GDR state -Shard the Qwen3.5 linear-attention/GDR path after the mixed-step and state-lifecycle contract is clear. +Shard the Qwen3.5 linear-attention/GDR path after P2a establishes the mixed-step and state-lifecycle contract. Expected work: @@ -169,6 +454,7 @@ Expected work: - 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` +- report matched Phase 1 TP2 versus P2b TP2 HBM/latency/throughput data before making a performance claim Non-negotiable invariant: @@ -176,6 +462,7 @@ Non-negotiable invariant: ## Follow-Ups +- Design and implement P2B sharded linear-attention/GDR state without weakening the completed P2A lifecycle and ID contracts. - 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/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 2acbb82d1..0638d222c 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -18,6 +18,9 @@ pub mod prefill_buffers; pub(crate) mod recurrent; pub(crate) mod recurrent_state; mod scheduler; +#[cfg(test)] +#[path = "../tests/common/model_fixture.rs"] +mod test_fixture; mod tp_executor; mod unified_forward; mod weights; @@ -61,6 +64,7 @@ pub mod runtime { pub use crate::executor::Qwen35Executor; pub use crate::executor::RequestId; pub use crate::scheduler::start_with_capacity; + pub use crate::tp_executor::DropExpectation; pub use crate::tp_executor::Qwen35TpExecutor; pub use crate::weights::Qwen35Model; } @@ -130,20 +134,6 @@ impl Qwen35LaunchOptions { } } -#[allow(clippy::needless_pass_by_value)] -fn launch_with_options_and_policy( - model_path: &Path, - options: Qwen35LaunchOptions, - scheduler_policy: Qwen35SchedulerPolicy, -) -> Result { - launch_with_options_policy_and_overlap( - model_path, - options, - scheduler_policy, - Qwen35DecodeOverlap::Off, - ) -} - #[allow(clippy::needless_pass_by_value)] pub fn launch_with_options_policy_and_overlap( model_path: &Path, diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs index d4259f44d..c2552c23c 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler.rs @@ -6,6 +6,9 @@ mod plan; +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; use std::sync::Arc; use std::sync::OnceLock; use std::sync::mpsc as std_mpsc; @@ -62,9 +65,11 @@ use crate::executor::PrefillResult; use crate::executor::RequestId; use crate::logprobs::snapshot_requested_logprobs; use crate::recurrent_state::RecurrentState; +use crate::tp_executor::DropExpectation; use crate::tp_executor::Qwen35TpExecutor; use crate::tp_executor::TpDecodeStepItem; use crate::tp_executor::TpPrefillChunkItem; +use crate::tp_executor::TpUnifiedPlan; use crate::weights::Qwen35Model; // ── Internal types ────────────────────────────────────────────────────── @@ -112,6 +117,162 @@ enum PrefillBackendState { Tp { request_id: RequestId }, } +struct TerminalRequest { + token_tx: TokenSink, + prompt_tokens: usize, + completion_tokens: usize, +} + +impl TerminalRequest { + fn send_error(self, message: &str) { + let _ = self.token_tx.send(TokenEvent::Error { + message: message.to_string(), + prompt_tokens: self.prompt_tokens, + completion_tokens: self.completion_tokens, + }); + } +} + +impl From for TerminalRequest { + fn from(req: SchedulerRequest) -> Self { + Self { + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + token_tx: req.token_tx, + } + } +} + +impl From for TerminalRequest { + fn from(req: ActiveRequest35) -> Self { + Self { + token_tx: req.token_tx, + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + } + } +} + +impl From for TerminalRequest { + fn from(req: PrefillingRequest35) -> Self { + req.req.into() + } +} + +struct PrefillCompletionRequest { + req: SchedulerRequest, + backend_state: PrefillBackendState, +} + +trait CompletionRequest { + fn token_tx(&self) -> &TokenSink; + fn into_terminal(self) -> TerminalRequest; +} + +impl CompletionRequest for ActiveRequest35 { + fn token_tx(&self) -> &TokenSink { + &self.token_tx + } + + fn into_terminal(self) -> TerminalRequest { + self.into() + } +} + +impl CompletionRequest for PrefillCompletionRequest { + fn token_tx(&self) -> &TokenSink { + &self.req.token_tx + } + + fn into_terminal(self) -> TerminalRequest { + self.req.into() + } +} + +struct CompletionCandidate { + request: R, + final_events: Vec, +} + +impl CompletionCandidate { + fn commit(self) { + for event in self.final_events { + let _ = self.request.token_tx().send(event); + } + } + + fn into_terminal(self) -> TerminalRequest { + self.request.into_terminal() + } +} + +struct FatalSchedulerError { + message: String, + transient: Vec, +} + +#[derive(Clone, Debug, PartialEq)] +struct PrefillArtifact { + token: u32, + logprob: Option, +} + +#[derive(Clone, Debug, PartialEq)] +struct DecodeArtifact { + token: u32, + logprob: Option, +} + +struct AlignedUnifiedArtifacts { + prefill: Vec>, + decode: Vec, +} + +enum PrefillStepArtifacts { + Single { + tokens: Vec, + logprobs: Vec>, + }, + Tp(Vec>), +} + +impl PrefillStepArtifacts { + fn final_artifact(&self, idx: usize) -> PrefillArtifact { + match self { + Self::Single { tokens, logprobs } => PrefillArtifact { + token: tokens[idx], + logprob: logprobs[idx].clone(), + }, + Self::Tp(artifacts) => artifacts[idx] + .clone() + .expect("validated TP final-prefill row must contain an artifact"), + } + } +} + +impl FatalSchedulerError { + fn new(message: impl Into) -> Self { + Self { + message: message.into(), + transient: Vec::new(), + } + } + + fn with_request(mut self, request: impl Into) -> Self { + self.transient.push(request.into()); + self + } + + fn with_requests(mut self, requests: I) -> Self + where + I: IntoIterator, + R: Into, + { + self.transient.extend(requests.into_iter().map(Into::into)); + self + } +} + pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; /// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` @@ -451,8 +612,7 @@ impl SingleGpuBackend { fn launch_async_prefill(&mut self, chunk: &mut ScheduledChunk) -> Result { let prefill_stream = self .prefill_stream - .as_ref() - .cloned() + .clone() .ok_or_else(|| anyhow::anyhow!("Qwen3.5 decode overlap is disabled"))?; // Request KV/recurrent state was allocated on the model stream. Order @@ -745,62 +905,53 @@ impl TpSchedulerBackend { &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(); + ) -> Result>> { + let items = tp_prefill_items(chunk)?; let result = self .executor .execute_prefill_chunks_with_seed(&items, sample_seed)?; align_prefill_results(chunk, &result) + .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) } 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::>()?; + ) -> Result> { + let items = tp_decode_items(active)?; let result = self.executor.execute_decode_items(&items, sample_seed)?; align_decode_results(active, &result) + .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) } - 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() - ); - } + fn execute_unified( + &self, + chunk: &ScheduledChunk, + active: &[ActiveRequest35], + decode_sample_seed: u64, + prefill_sample_seed: u64, + ) -> Result { + let plan = TpUnifiedPlan { + prefill: tp_prefill_items(chunk)?, + decode: tp_decode_items(active)?, + prefill_sample_seed, + decode_sample_seed, + }; + let result = self.executor.execute_unified(&plan)?; + let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { + self.executor + .poison_artifact_contract("unified prefill", &err) + })?; + let decode = align_decode_results(active, &result.decode).map_err(|err| { + self.executor + .poison_artifact_contract("unified decode", &err) + })?; + Ok(AlignedUnifiedArtifacts { prefill, decode }) + } + + fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { + self.executor.drop_request(request_id, expectation) } } @@ -856,10 +1007,6 @@ impl SchedulerBackend { } } - 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), @@ -877,73 +1024,174 @@ fn pages_needed(token_count: usize, page_size: usize) -> usize { token_count.div_ceil(page_size) } +fn tp_prefill_items(chunk: &ScheduledChunk) -> Result> { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + anyhow::bail!("TP prefill received single-GPU chunk state"); + }; + anyhow::ensure!( + chunk.reqs.len() == request_ids.len() + && chunk.reqs.len() == chunk.windows.len() + && chunk.reqs.len() == chunk.ends.len(), + "Qwen3.5 TP scheduled prefill vectors are misaligned" + ); + Ok(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()) +} + +fn tp_decode_items(active: &[ActiveRequest35]) -> Result> { + 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() +} + fn align_prefill_results( chunk: &ScheduledChunk, result: &PrefillResult, -) -> Result<(Vec, Vec>)> { +) -> Result>> { 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()]; + anyhow::ensure!( + request_ids.len() == chunk.reqs.len() && chunk.ends.len() == chunk.reqs.len(), + "Qwen3.5 TP prefill alignment vectors are misaligned" + ); + let expected: HashSet = request_ids + .iter() + .zip(&chunk.reqs) + .zip(&chunk.ends) + .filter_map(|((&request_id, req), &end)| { + (end == req.prompt_tokens.len()).then_some(request_id) + }) + .collect(); + let mut by_id = HashMap::with_capacity(result.requests.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].clone_from(first_token_logprob); + anyhow::ensure!( + expected.contains(request_id), + "Qwen3.5 TP prefill returned unknown or non-final request id {}", + request_id.get() + ); + let artifact = PrefillArtifact { + token: *first_token, + logprob: first_token_logprob.clone(), + }; + anyhow::ensure!( + by_id.insert(*request_id, artifact).is_none(), + "Qwen3.5 TP prefill returned duplicate request id {}", + request_id.get() + ); } - Ok((tokens, logprobs)) + anyhow::ensure!( + by_id.len() == expected.len(), + "Qwen3.5 TP prefill result is missing final request IDs" + ); + + request_ids + .iter() + .zip(&chunk.reqs) + .zip(&chunk.ends) + .map(|((&request_id, req), &end)| { + if end == req.prompt_tokens.len() { + by_id.remove(&request_id).map(Some).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP prefill result is missing final request id {}", + request_id.get() + ) + }) + } else { + Ok(None) + } + }) + .collect() } fn align_decode_results( active: &[ActiveRequest35], result: &DecodeResult, -) -> Result<(Vec, Vec>)> { +) -> Result> { + let expected: Vec = active + .iter() + .map(|active_req| { + let ActiveBackendState::Tp { request_id } = active_req.backend_state else { + anyhow::bail!("align_decode_results requires TP active state"); + }; + Ok(request_id) + }) + .collect::>()?; + let expected_set: HashSet<_> = expected.iter().copied().collect(); anyhow::ensure!( - active.len() == result.requests.len(), - "Qwen3.5 TP decode result row count mismatch: active={}, result={}", - active.len(), - result.requests.len() + expected_set.len() == expected.len(), + "Qwen3.5 TP active decode IDs contain duplicates" ); - 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 mut by_id = HashMap::with_capacity(result.requests.len()); + for DecodeRequestResult { + request_id, + token, + logprob, + } in &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_set.contains(request_id), + "Qwen3.5 TP decode returned unknown request id {}", + request_id.get() + ); + let artifact = DecodeArtifact { + token: *token, + logprob: logprob.clone(), }; anyhow::ensure!( - *expected == *request_id, - "Qwen3.5 TP decode result request id mismatch: expected {}, got {}", - expected.get(), + by_id.insert(*request_id, artifact).is_none(), + "Qwen3.5 TP decode returned duplicate request id {}", request_id.get() ); - tokens.push(*token); - logprobs.push(logprob.clone()); } - Ok((tokens, logprobs)) + expected + .into_iter() + .map(|request_id| { + by_id.remove(&request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode result is missing request id {}", + request_id.get() + ) + }) + }) + .collect() +} + +fn split_decode_artifacts(artifacts: &[DecodeArtifact]) -> (Vec, Vec>) { + artifacts + .iter() + .map(|artifact| (artifact.token, artifact.logprob.clone())) + .unzip() } fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { @@ -996,23 +1244,136 @@ fn publish_load( num_waiting_reqs: usize, ) { let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + let (num_running_reqs, num_waiting_reqs) = + logical_load_counts(active, prefilling, inflight_prefill_reqs, num_waiting_reqs); load_tx.send_replace(SchedulerMetrics { kv_used_blocks: kv_total_blocks .saturating_sub(backend.available_pages(active, prefilling) as u64), kv_total_blocks, - num_running_reqs: (active.len() + prefilling.len() + inflight_prefill_reqs) as u64, - num_waiting_reqs: num_waiting_reqs as u64, + num_running_reqs, + num_waiting_reqs, spec_decode: None, }); } -fn should_block_on_submit( - active_empty: bool, - prefilling_empty: bool, - pending_empty: bool, - inflight_prefill: bool, -) -> bool { - active_empty && prefilling_empty && pending_empty && !inflight_prefill +fn logical_load_counts( + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + inflight_prefill_reqs: usize, + num_waiting_reqs: usize, +) -> (u64, u64) { + ( + (active.len() + prefilling.len() + inflight_prefill_reqs) as u64, + num_waiting_reqs as u64, + ) +} + +fn should_block_on_submit(owned_work_empty: bool, inflight_prefill: bool) -> bool { + owned_work_empty && !inflight_prefill +} + +fn terminal_scheduler_shutdown( + submit_rx: &mut mpsc::UnboundedReceiver, + load_tx: &watch::Sender, + kv_total_blocks: u64, + active: Vec, + prefilling: Vec, + pending: Vec, + deferred: Vec, + inflight_prefill: Option, + failure: FatalSchedulerError, +) { + submit_rx.close(); + + let mut requests = failure.transient; + requests.extend(active.into_iter().map(Into::into)); + requests.extend(prefilling.into_iter().map(Into::into)); + requests.extend(pending.into_iter().map(Into::into)); + requests.extend(deferred.into_iter().map(Into::into)); + if let Some(InflightPrefill { output, chunk, .. }) = inflight_prefill { + // The stream must drain before the chunk's KV/recurrent/conv state is + // released or transferred into terminal request ownership. + drop(output); + requests.extend(chunk.reqs.into_iter().map(Into::into)); + } + while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { + requests.push(req.into()); + } + + warn!( + "Qwen3.5 TP scheduler terminating after replica failure: {}", + failure.message + ); + for request in requests { + request.send_error(&failure.message); + } + load_tx.send_replace(SchedulerMetrics { + kv_used_blocks: 0, + kv_total_blocks, + num_running_reqs: 0, + num_waiting_reqs: 0, + spec_decode: None, + }); +} + +fn prune_closed_requests( + backend: &mut B, + active: &mut Vec, + prefilling: &mut Vec, + pending: &mut Vec, +) -> std::result::Result<(), FatalSchedulerError> +where + B: DecodeDispatchBackend + PrefillPromoteBackend, +{ + pending.retain(|req| !req.token_tx.is_closed()); + + for idx in (0..active.len()).rev() { + if active[idx].token_tx.is_closed() { + debug!( + "request pruned before scheduling: request_id={:?} phase=decode tokens_generated={}", + active[idx].request_id, active[idx].generated_count + ); + let removed = backend.take_active_request(active, idx); + if let Err(err) = backend.drop_active_state(&removed.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); + } + } + } + + for idx in (0..prefilling.len()).rev() { + if prefilling[idx].req.token_tx.is_closed() { + let removed = prefilling.remove(idx); + debug!( + "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", + removed.req.request_id, removed.cursor + ); + let expectation = if removed.cursor == 0 { + DropExpectation::MustBeAbsent + } else { + DropExpectation::MustExist + }; + if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); + } + } + } + Ok(()) +} + +const UNSUPPORTED_ECHO_MESSAGE: &str = "echo=true is unsupported by the Qwen3.5 serving contract"; + +fn reject_unsupported_echo(pending: &mut Vec) { + pending.retain(|req| { + if !req.echo { + return true; + } + let _ = req.token_tx.send(TokenEvent::Rejected { + message: UNSUPPORTED_ECHO_MESSAGE.to_string(), + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + }); + false + }); } #[allow(clippy::needless_pass_by_value)] @@ -1034,20 +1395,6 @@ fn scheduler_loop( info!("scheduler ready (max_batch={})", max_batch); loop { - // Publish the settled state between scheduler steps. If the prior step - // retired its final requests, their KV pages have already returned via - // RAII, so this snapshot reaches idle before the channel blocks below. - publish_load( - &load_tx, - &backend, - &active, - &prefilling, - inflight_prefill - .as_ref() - .map_or(0, |prefill| prefill.chunk.reqs.len()), - deferred.len(), - ); - if inflight_prefill .as_mut() .is_some_and(|prefill| prefill.output.is_ready()) @@ -1061,7 +1408,7 @@ fn scheduler_loop( }); let decode_n = active.len(); let step_start = itl_debug_enabled().then(Instant::now); - finish_async_prefill( + let finish_result = finish_async_prefill( &mut backend, &mut active, &mut prefilling, @@ -1076,20 +1423,71 @@ fn scheduler_loop( prefill_reqs, decode_n, ); + if let Err(failure) = finish_result { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + Vec::new(), + deferred, + inflight_prefill.take(), + failure, + ); + return; + } } - // 1. Drain all pending requests (deferred from last iteration + channel) + // 1. Merge deferred work with every submission currently available. let mut pending = std::mem::take(&mut deferred); while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { pending.push(req); } - // 2. Nothing in flight (no decode, no in-progress prefill) and nothing - // pending → block until a request arrives. + // 2. Remove closed work before metrics, admission, or planning. Active + // and prefilling cleanup goes through the backend's normal retirement + // paths so graph slots and TP request state are released consistently. + if let Err(failure) = + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) + { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + pending, + deferred, + inflight_prefill.take(), + failure, + ); + return; + } + reject_unsupported_echo(&mut pending); + + // 3. Publish the settled post-prune state. Requests accepted from the + // channel are waiting until admission below; closed requests never + // appear in this snapshot or consume its KV/slot accounting. + publish_load( + &load_tx, + &backend, + &active, + &prefilling, + inflight_prefill + .as_ref() + .map_or(0, |prefill| prefill.chunk.reqs.len()), + pending.len(), + ); + + // 4. Nothing in flight and nothing pending: the idle snapshot above is + // already visible, so block until work arrives. Drain and prune again + // after wakeup because the first request may already be closed and more + // submissions may have raced with the blocking receive. if should_block_on_submit( - active.is_empty(), - prefilling.is_empty(), - pending.is_empty(), + active.is_empty() && prefilling.is_empty() && pending.is_empty(), inflight_prefill.is_some(), ) { if let Some((req, _kv_prefix)) = submit_rx.blocking_recv() { @@ -1101,6 +1499,28 @@ fn scheduler_loop( while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { pending.push(req); } + if let Err(failure) = + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) + { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + pending, + deferred, + inflight_prefill.take(), + failure, + ); + return; + } + reject_unsupported_echo(&mut pending); + publish_load(&load_tx, &backend, &active, &prefilling, 0, pending.len()); + if pending.is_empty() { + continue; + } } // One async prefill owns its scheduled request state. Do not admit or @@ -1117,8 +1537,8 @@ fn scheduler_loop( ) }); let itl_decode_n = active.len(); - let itl_plan_kind = if active.is_empty() { - finish_async_prefill( + let (itl_plan_kind, step_result) = if active.is_empty() { + let result = finish_async_prefill( &mut backend, &mut active, &mut prefilling, @@ -1126,10 +1546,10 @@ fn scheduler_loop( .take() .expect("async prefill must be present before blocking wait"), ); - "overlap_wait" + ("overlap_wait", result) } else { - decode_step(&mut backend, &mut active, &mut rng); - "overlap_decode" + let result = decode_step(&mut backend, &mut active, &mut rng); + ("overlap_decode", result) }; log_itl_step( itl_step_start, @@ -1138,10 +1558,25 @@ fn scheduler_loop( itl_prefill_reqs, itl_decode_n, ); + if let Err(failure) = step_result { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + Vec::new(), + deferred, + inflight_prefill.take(), + failure, + ); + return; + } continue; } - // 3. Admit new prompts. In-flight prefills reserve their promotion slot + // 5. Admit new prompts. In-flight prefills reserve their promotion slot // and future KV growth, so shrink the slot/page budgets accordingly let active_budget: Vec = active .iter() @@ -1181,7 +1616,7 @@ fn scheduler_loop( send_rejection(rejected, *reason); } - // 4. Move freshly admitted prompts into the chunked-prefill queue. + // 6. Move freshly admitted prompts into the chunked-prefill queue. for req in admission.pending { debug!( "request admitted: request_id={:?} prompt_len={} max_tokens={}", @@ -1209,7 +1644,7 @@ fn scheduler_loop( deferred = admission.deferred; - // 5. Choose this tick's prefill budget, take that chunk off the front of + // 7. Choose this tick's prefill budget, take that chunk off the front of // the queue, then dispatch by plan. Auto can return 0 for a short // decode-priority tick; the next iteration reconsiders the same FIFO // prefill without reordering it. @@ -1240,11 +1675,7 @@ fn scheduler_loop( let itl_prefill_tokens: usize = scheduled.iter().map(|p| p.step_chunk).sum(); let itl_prefill_reqs = scheduled.len(); let itl_decode_n = active.len(); - let plan = if backend.is_tp() { - build_eager_only_plan(!active.is_empty(), scheduled) - } else { - plan::build_next_plan(!active.is_empty(), scheduled) - }; + let plan = plan::build_next_plan(!active.is_empty(), scheduled); if let Some(plan) = plan { let itl_plan_kind = match &plan { ExecutionPlan::Unified { .. } if matches!(&backend, SchedulerBackend::Single(single) if single.overlap_enabled()) => { @@ -1255,7 +1686,7 @@ fn scheduler_loop( ExecutionPlan::Decode => "decode", }; let itl_step_start = itl_debug.then(Instant::now); - match plan { + let step_result = match plan { ExecutionPlan::Unified { pending } => { if matches!(&backend, SchedulerBackend::Single(single) if single.overlap_enabled()) { @@ -1265,7 +1696,7 @@ fn scheduler_loop( pending, &mut inflight_prefill, &mut rng, - ); + ) } else { unified_step_sched( &mut backend, @@ -1273,7 +1704,7 @@ fn scheduler_loop( pending, &mut prefilling, &mut rng, - ); + ) } } ExecutionPlan::Prefill { pending } => prefill_batch( @@ -1283,10 +1714,8 @@ fn scheduler_loop( &mut prefilling, &mut rng, ), - ExecutionPlan::Decode => { - decode_step(&mut backend, &mut active, &mut rng); - } - } + ExecutionPlan::Decode => decode_step(&mut backend, &mut active, &mut rng), + }; log_itl_step( itl_step_start, itl_plan_kind, @@ -1294,20 +1723,25 @@ fn scheduler_loop( itl_prefill_reqs, itl_decode_n, ); + if let Err(failure) = step_result { + let kv_total_blocks = backend.capacity_pages_for_requests() as u64; + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + kv_total_blocks, + active, + prefilling, + Vec::new(), + deferred, + inflight_prefill.take(), + failure, + ); + return; + } } } } -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!( @@ -1339,10 +1773,10 @@ fn prefill_batch( scheduled: Vec, prefilling: &mut Vec, rng: &mut StdRng, -) { +) -> std::result::Result<(), FatalSchedulerError> { let mut chunk = ScheduledChunk::from(scheduled); let sample_seed = rand::RngExt::random(rng); - let (tokens, logprobs_vec) = match backend { + let artifacts = match backend { SchedulerBackend::Single(single) => { // Scope the borrows of `chunk` to the executor call so the error path can // move `chunk` into `fail_chunk`. @@ -1351,31 +1785,29 @@ fn prefill_batch( Err(e) => { warn!("batch prefill failed: {e}"); fail_chunk(chunk, &e.to_string()); - return; + return Ok(()); } }; let prefill_sample_seed = rand::RngExt::random(rng); match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { - Ok(v) => v, + Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, Err(e) => { warn!("prefill sampling failed: {e}"); fail_chunk(chunk, &e.to_string()); - return; + return Ok(()); } } } SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { - Ok(v) => v, + Ok(v) => PrefillStepArtifacts::Tp(v), Err(e) => { warn!("TP prefill chunk failed: {e}"); - drop_tp_chunk_state(tp, &chunk); - fail_chunk(chunk, &e.to_string()); - return; + return Err(FatalSchedulerError::new(e.to_string()).with_requests(chunk.reqs)); } }, }; - promote_or_requeue(backend, active, prefilling, chunk, &tokens, &logprobs_vec); + promote_or_requeue(backend, active, prefilling, chunk, &artifacts) } fn launch_overlap_step( @@ -1384,7 +1816,7 @@ fn launch_overlap_step( scheduled: Vec, inflight_prefill: &mut Option, rng: &mut StdRng, -) { +) -> std::result::Result<(), FatalSchedulerError> { debug_assert!(inflight_prefill.is_none()); let mut chunk = ScheduledChunk::from(scheduled); let decode_seed = rand::RngExt::random(rng); @@ -1406,7 +1838,7 @@ fn launch_overlap_step( fail_chunk(chunk, &err.to_string()); } } - decode_step_with_seed(backend, active, decode_seed); + decode_step_with_seed(backend, active, decode_seed) } fn finish_async_prefill( @@ -1414,7 +1846,7 @@ fn finish_async_prefill( active: &mut Vec, prefilling: &mut Vec, inflight: InflightPrefill, -) { +) -> std::result::Result<(), FatalSchedulerError> { let InflightPrefill { chunk, output, @@ -1429,10 +1861,11 @@ fn finish_async_prefill( Err(err) => { warn!("async prefill sampling failed: {err}"); fail_chunk(chunk, &err.to_string()); - return; + return Ok(()); } }; - promote_or_requeue(single, active, prefilling, chunk, &tokens, &logprobs); + let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(single, active, prefilling, chunk, &artifacts) } // ── Unified step (prefill chunk + decode in one forward pass) ────────────── @@ -1443,22 +1876,42 @@ fn unified_step_sched( scheduled: Vec, prefilling: &mut Vec, 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, - }); +) -> std::result::Result<(), FatalSchedulerError> { + let mut chunk = ScheduledChunk::from(scheduled); + if matches!(backend, SchedulerBackend::Tp(_)) { + // Preserve the established scheduler RNG order: decode seed first, + // prefill seed second. Workers execute the forwards in the opposite + // (prefill-then-decode) order using these preselected seeds. + let decode_sample_seed = rand::RngExt::random(rng); + let prefill_sample_seed = rand::RngExt::random(rng); + let result = { + let SchedulerBackend::Tp(tp) = backend else { + unreachable!() + }; + tp.execute_unified(&chunk, active, decode_sample_seed, prefill_sample_seed) + }; + let artifacts = match result { + Ok(artifacts) => artifacts, + Err(err) => { + warn!("TP unified step failed: {err}"); + return Err(FatalSchedulerError::new(err.to_string()).with_requests(chunk.reqs)); + } + }; + + let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); + if let Err(failure) = + dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs) + { + return Err(failure.with_requests(chunk.reqs)); } - fail_chunk(chunk, message); - return; + + let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); + return promote_or_requeue(backend, active, prefilling, chunk, &prefill); + } + + let SchedulerBackend::Single(backend) = backend else { + unreachable!() }; - 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 = backend.unified_step(&mut chunk, active); @@ -1475,7 +1928,7 @@ fn unified_step_sched( }); } fail_chunk(chunk, &message); - return; + return Ok(()); } }; let decode_seed = rand::RngExt::random(rng); @@ -1484,24 +1937,24 @@ 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(backend, active, decode_seed); + process_decode_logits(backend, active, decode_seed)?; } let prefill_logits = output .prefill_logits .as_ref() .expect("scheduled prefill chunk must return prefill logits"); - let (tokens, logprobs_vec) = + let (tokens, logprobs) = match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, prefill_seed) { Ok(v) => v, Err(e) => { warn!("unified prefill sampling failed: {e}"); fail_chunk(chunk, &e.to_string()); - return; + return Ok(()); } }; - - promote_or_requeue(backend, active, prefilling, chunk, &tokens, &logprobs_vec); + let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(backend, active, prefilling, chunk, &prefill) } // ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── @@ -1510,7 +1963,7 @@ fn decode_step( backend: &mut SchedulerBackend, active: &mut Vec, rng: &mut StdRng, -) { +) -> std::result::Result<(), FatalSchedulerError> { // Preserve the historical scheduler RNG sequence: TP consumes the first // seed, while single-GPU decode consumed a second seed inside sampling. let first_seed = rand::RngExt::random(rng); @@ -1519,14 +1972,14 @@ fn decode_step( } else { first_seed }; - decode_step_with_seed(backend, active, sample_seed); + decode_step_with_seed(backend, active, sample_seed) } fn decode_step_with_seed( backend: &mut SchedulerBackend, active: &mut Vec, sample_seed: u64, -) { +) -> std::result::Result<(), FatalSchedulerError> { let (tokens, logprobs_vec) = match backend { SchedulerBackend::Single(single) => { if let Err(e) = single.decode_graph(active) { @@ -1539,7 +1992,7 @@ fn decode_step_with_seed( completion_tokens: req.generated_count, }); } - return; + return Ok(()); } // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) match single.sample_decode_logits(active, sample_seed) { @@ -1554,32 +2007,20 @@ fn decode_step_with_seed( completion_tokens: req.generated_count, }); } - return; + return Ok(()); } } } SchedulerBackend::Tp(tp) => match tp.execute_decode(active, sample_seed) { - Ok(v) => v, + Ok(v) => split_decode_artifacts(&v), Err(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, - completion_tokens: req.generated_count, - }); - } - return; + return Err(FatalSchedulerError::new(e.to_string())); } }, }; - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) } /// Process decode logits from unified step: sample, extract logprobs, dispatch. @@ -1587,7 +2028,7 @@ fn process_decode_logits( backend: &mut SingleGpuBackend, active: &mut Vec, sample_seed: u64, -) { +) -> std::result::Result<(), FatalSchedulerError> { let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { Ok(v) => v, Err(e) => { @@ -1600,11 +2041,11 @@ fn process_decode_logits( completion_tokens: req.generated_count, }); } - return; + return Ok(()); } }; - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec); + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) } /// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. @@ -1616,7 +2057,13 @@ fn dispatch_decode_tokens( active: &mut Vec, tokens: &[u32], logprobs: &[Option], -) { +) -> std::result::Result<(), FatalSchedulerError> { + enum Retirement { + Completion(Vec), + CleanupOnly, + Disconnected, + } + let n = active.len(); let mut to_retire = Vec::new(); @@ -1637,12 +2084,17 @@ fn dispatch_decode_tokens( req.generated_count, FinishReason::Stop ); - let _ = req.token_tx.send(TokenEvent::Finished { + let event = TokenEvent::Finished { finish_reason: FinishReason::Stop, prompt_tokens: req.prompt_len, completion_tokens: req.generated_count, - }); - to_retire.push(i); + }; + if backend.completion_requires_drop_ack() { + to_retire.push((i, Retirement::Completion(vec![event]))); + } else { + let _ = req.token_tx.send(event); + to_retire.push((i, Retirement::CleanupOnly)); + } } else if at_limit { debug!( "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", @@ -1651,13 +2103,22 @@ fn dispatch_decode_tokens( req.generated_count, FinishReason::Length ); - let _ = req.token_tx.send(TokenEvent::Token { id: token, logprob }); - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - to_retire.push(i); + let events = vec![ + TokenEvent::Token { id: token, logprob }, + TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }, + ]; + if backend.completion_requires_drop_ack() { + to_retire.push((i, Retirement::Completion(events))); + } else { + for event in events { + let _ = req.token_tx.send(event); + } + to_retire.push((i, Retirement::CleanupOnly)); + } } else if req .token_tx .send(TokenEvent::Token { id: token, logprob }) @@ -1667,21 +2128,46 @@ fn dispatch_decode_tokens( "request dropped: client disconnected: request_id={:?} tokens_generated={}", req.request_id, req.generated_count ); - to_retire.push(i); + to_retire.push((i, Retirement::Disconnected)); } else { req.last_token = token; } } // Remove in reverse order so compact_slot indices stay valid - for &i in to_retire.iter().rev() { - backend.retire_request(active, i); + for (i, retirement) in to_retire.into_iter().rev() { + let request = backend.take_active_request(active, i); + match retirement { + Retirement::Completion(final_events) => { + let candidate = CompletionCandidate { + request, + final_events, + }; + if let Err(err) = backend.drop_active_state(&candidate.request.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()) + .with_request(candidate.into_terminal())); + } + candidate.commit(); + } + Retirement::CleanupOnly | Retirement::Disconnected => { + if let Err(err) = backend.drop_active_state(&request.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(request)); + } + } + } } + Ok(()) } trait DecodeDispatchBackend { fn is_stop_token(&self, token: u32) -> bool; - fn retire_request(&mut self, active: &mut Vec, idx: usize); + fn completion_requires_drop_ack(&self) -> bool; + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35; + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()>; } impl DecodeDispatchBackend for SingleGpuBackend { @@ -1689,8 +2175,20 @@ impl DecodeDispatchBackend for SingleGpuBackend { self.is_stop_token(token) } - fn retire_request(&mut self, active: &mut Vec, idx: usize) { - compact_single_slot(self, active, idx); + fn completion_requires_drop_ack(&self) -> bool { + false + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + compact_single_slot(self, active, idx) + } + + fn drop_active_state(&mut self, _state: &ActiveBackendState) -> Result<()> { + Ok(()) } } @@ -1699,15 +2197,28 @@ impl DecodeDispatchBackend for SchedulerBackend { self.is_stop_token(token) } - fn retire_request(&mut self, active: &mut Vec, idx: usize) { + fn completion_requires_drop_ack(&self) -> bool { + matches!(self, SchedulerBackend::Tp(_)) + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { 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); - } + SchedulerBackend::Tp(_) => active.swap_remove(idx), + } + } + + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { + match (self, state) { + (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), + (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id }) => { + backend.drop_request(*request_id, DropExpectation::MustExist) } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), } } } @@ -1721,13 +2232,14 @@ fn compact_single_slot( backend: &mut SingleGpuBackend, active: &mut Vec, idx: usize, -) { +) -> ActiveRequest35 { let compaction = compaction_after_retire(active.len(), idx); - active.swap_remove(idx); + let removed = active.swap_remove(idx); if let Some(compaction) = compaction { backend.compact_slot(active, compaction); } + removed } // ── Chunked-prefill helpers ──────────────────────────────────────────────── @@ -1836,27 +2348,17 @@ 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`. +/// `artifacts` are indexed by request order in `chunk`. fn promote_or_requeue( backend: &mut impl PrefillPromoteBackend, active: &mut Vec, prefilling: &mut Vec, chunk: ScheduledChunk, - tokens: &[u32], - logprobs: &[Option], -) { + artifacts: &PrefillStepArtifacts, +) -> std::result::Result<(), FatalSchedulerError> { let ScheduledChunk { reqs, backend_state, @@ -1865,10 +2367,15 @@ fn promote_or_requeue( } = chunk; let mut still_prefilling: Vec = Vec::new(); let backend_states = split_scheduled_backend_state(backend_state); + let mut entries: VecDeque<_> = reqs + .into_iter() + .zip(backend_states) + .zip(ends) + .enumerate() + .map(|(i, ((req, backend_state), end))| (i, req, backend_state, end)) + .collect(); - for (i, ((req, backend_state), end)) in - reqs.into_iter().zip(backend_states).zip(ends).enumerate() - { + while let Some((i, req, backend_state, end)) = entries.pop_front() { // Not finished: re-queue with the advanced cursor if end < req.prompt_tokens.len() { still_prefilling.push(PrefillingRequest35 { @@ -1881,16 +2388,9 @@ fn promote_or_requeue( } let prompt_len = req.prompt_tokens.len(); - let first_token = tokens[i]; - let logprob = logprobs[i].clone(); - - if req.echo { - let echo_logprobs = vec![None; req.prompt_tokens.len()]; - let _ = req.token_tx.send(TokenEvent::PromptTokens { - ids: req.prompt_tokens.clone(), - logprobs: echo_logprobs, - }); - } + let artifact = artifacts.final_artifact(i); + let first_token = artifact.token; + let logprob = artifact.logprob; if !req.params.ignore_eos && backend.is_stop_token(first_token) { debug!( @@ -1900,12 +2400,61 @@ fn promote_or_requeue( 0, FinishReason::Stop ); - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }); - backend.drop_prefill_state(backend_state); + let candidate = CompletionCandidate { + request: PrefillCompletionRequest { req, backend_state }, + final_events: vec![TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: prompt_len, + completion_tokens: 0, + }], + }; + if let Err(err) = backend + .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + candidate.into_terminal(), + still_prefilling, + entries, + )); + } + candidate.commit(); + continue; + } + + if req.max_tokens <= 1 { + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + req.request_id, + prompt_len, + 1, + FinishReason::Length + ); + let candidate = CompletionCandidate { + request: PrefillCompletionRequest { req, backend_state }, + final_events: vec![ + TokenEvent::Token { + id: first_token, + logprob, + }, + TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: prompt_len, + completion_tokens: 1, + }, + ], + }; + if let Err(err) = backend + .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + candidate.into_terminal(), + still_prefilling, + entries, + )); + } + candidate.commit(); continue; } @@ -1921,24 +2470,17 @@ fn promote_or_requeue( "request dropped: client disconnected: request_id={:?} tokens_generated={}", req.request_id, 0 ); - backend.drop_prefill_state(backend_state); - continue; - } - - if req.max_tokens <= 1 { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 1, - FinishReason::Length - ); - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: prompt_len, - completion_tokens: 1, - }); - backend.drop_prefill_state(backend_state); + let removed = PrefillCompletionRequest { req, backend_state }; + if let Err(err) = + backend.drop_prefill_state(&removed.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + removed.into_terminal(), + still_prefilling, + entries, + )); + } continue; } @@ -1957,6 +2499,19 @@ fn promote_or_requeue( } prefilling.splice(0..0, still_prefilling); + Ok(()) +} + +fn prefill_lifecycle_failure( + message: String, + current: TerminalRequest, + still_prefilling: Vec, + remaining: VecDeque<(usize, SchedulerRequest, PrefillBackendState, usize)>, +) -> FatalSchedulerError { + FatalSchedulerError::new(message) + .with_request(current) + .with_requests(still_prefilling) + .with_requests(remaining.into_iter().map(|(_, req, _, _)| req)) } trait PrefillPromoteBackend { @@ -1966,7 +2521,11 @@ trait PrefillPromoteBackend { active_len: usize, state: PrefillBackendState, ) -> ActiveBackendState; - fn drop_prefill_state(&mut self, state: PrefillBackendState); + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()>; } impl PrefillPromoteBackend for SingleGpuBackend { @@ -1992,7 +2551,13 @@ impl PrefillPromoteBackend for SingleGpuBackend { } } - fn drop_prefill_state(&mut self, _state: PrefillBackendState) {} + fn drop_prefill_state( + &mut self, + _state: &PrefillBackendState, + _expectation: DropExpectation, + ) -> Result<()> { + Ok(()) + } } impl PrefillPromoteBackend for SchedulerBackend { @@ -2024,11 +2589,17 @@ impl PrefillPromoteBackend for SchedulerBackend { } } - 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 drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()> { + match (self, state) { + (SchedulerBackend::Single(_), PrefillBackendState::Single { .. }) => Ok(()), + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + backend.drop_request(*request_id, expectation) + } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during prefill drop"), } } } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index c952fb275..39a521517 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -1,10 +1,803 @@ use std::path::Path; +use std::sync::Arc; +use std::sync::Barrier; +use std::time::Duration; +use std::time::Instant; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::EpBackend; use super::*; +fn test_request(request_id: &str, token_tx: TokenSink) -> SchedulerRequest { + test_request_with_shape(request_id, token_tx, vec![1], 1) +} + +fn test_request_with_shape( + request_id: &str, + token_tx: TokenSink, + prompt_tokens: Vec, + max_tokens: usize, +) -> SchedulerRequest { + SchedulerRequest { + trace_parent: None, + request_id: Some(request_id.to_string()), + queued_at_unix_s: None, + data_parallel_rank: None, + prompt_tokens, + params: SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + lora_adapter: None, + kv_transfer_params: None, + token_tx, + logprobs: 0, + echo: false, + } +} + +fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRequest35 { + ActiveRequest35 { + request_id: Some(label.to_string()), + token_tx, + backend_state: ActiveBackendState::Tp { + request_id: RequestId::new(request_id), + }, + last_token: 1, + generated_count: 1, + max_tokens: 8, + prompt_len: 1, + params: SamplingParams::default(), + logprobs: 0, + } +} + +fn prefilling_request(request_id: u64, label: &str, token_tx: TokenSink) -> PrefillingRequest35 { + PrefillingRequest35 { + req: test_request(label, token_tx), + backend_state: PrefillBackendState::Tp { + request_id: RequestId::new(request_id), + }, + cursor: 0, + step_chunk: 0, + } +} + +#[derive(Default)] +struct PruneTestBackend { + retired_active: Vec, + dropped_prefilling: Vec<(RequestId, DropExpectation)>, +} + +impl DecodeDispatchBackend for PruneTestBackend { + fn is_stop_token(&self, _token: u32) -> bool { + false + } + + fn completion_requires_drop_ack(&self) -> bool { + true + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + active.swap_remove(idx) + } + + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { + let ActiveBackendState::Tp { request_id } = state else { + panic!("prune test expected TP active state"); + }; + self.retired_active.push(*request_id); + Ok(()) + } +} + +impl PrefillPromoteBackend for PruneTestBackend { + fn is_stop_token(&self, _token: u32) -> bool { + false + } + + fn promote_prefill_state( + &mut self, + _active_len: usize, + _state: PrefillBackendState, + ) -> ActiveBackendState { + panic!("prune test must not promote prefill state") + } + + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()> { + let PrefillBackendState::Tp { request_id } = state else { + panic!("prune test expected TP prefill state"); + }; + self.dropped_prefilling.push((*request_id, expectation)); + Ok(()) + } +} + +struct LifecycleTestBackend { + stop_token: Option, + active_completion_requires_drop_ack: bool, + fail_active_drop: bool, + fail_prefill_drop: bool, + active_drops: Vec, + active_events_before_drop: Vec, + prefill_drops: Vec<(RequestId, DropExpectation)>, + observer: Option, +} + +impl LifecycleTestBackend { + fn new( + stop_token: Option, + observer: pegainfer_frontend::engine::TokenStreamReceiver, + ) -> Self { + Self { + stop_token, + active_completion_requires_drop_ack: true, + fail_active_drop: false, + fail_prefill_drop: false, + active_drops: Vec::new(), + active_events_before_drop: Vec::new(), + prefill_drops: Vec::new(), + observer: Some(observer), + } + } + + fn assert_no_completion_published(&mut self) { + let Some(observer) = &mut self.observer else { + return; + }; + assert!(matches!( + observer.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + } +} + +impl DecodeDispatchBackend for LifecycleTestBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.stop_token == Some(token) + } + + fn completion_requires_drop_ack(&self) -> bool { + self.active_completion_requires_drop_ack + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + active.swap_remove(idx) + } + + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { + let ActiveBackendState::Tp { request_id } = state else { + panic!("lifecycle test expected TP active state"); + }; + if self.active_completion_requires_drop_ack { + self.assert_no_completion_published(); + } else if let Some(observer) = &mut self.observer { + while let Ok((_, event)) = observer.try_recv() { + self.active_events_before_drop.push(event); + } + } + self.active_drops.push(*request_id); + anyhow::ensure!(!self.fail_active_drop, "injected active drop failure"); + Ok(()) + } +} + +impl PrefillPromoteBackend for LifecycleTestBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.stop_token == Some(token) + } + + fn promote_prefill_state( + &mut self, + _active_len: usize, + _state: PrefillBackendState, + ) -> ActiveBackendState { + panic!("completion lifecycle test must not promote prefill state") + } + + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()> { + let PrefillBackendState::Tp { request_id } = state else { + panic!("lifecycle test expected TP prefill state"); + }; + self.assert_no_completion_published(); + self.prefill_drops.push((*request_id, expectation)); + anyhow::ensure!(!self.fail_prefill_drop, "injected prefill drop failure"); + Ok(()) + } +} + +fn next_event( + rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, + description: &str, +) -> TokenEvent { + rx.blocking_recv() + .unwrap_or_else(|| panic!("{description} channel closed before event")) + .1 +} + +fn assert_no_more_events(rx: &mut pegainfer_frontend::engine::TokenStreamReceiver) { + assert!( + rx.try_recv().is_err(), + "request received more than one terminal event" + ); +} + +#[test] +fn closed_pending_work_is_pruned_before_admission() { + let (closed_sink, closed_rx) = TokenSink::standalone(); + drop(closed_rx); + let (open_sink, _open_rx) = TokenSink::standalone(); + let mut pending = vec![ + test_request("closed", closed_sink), + test_request("open", open_sink), + ]; + let mut active = Vec::new(); + let mut prefilling = Vec::new(); + let mut backend = PruneTestBackend::default(); + + assert!( + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() + ); + + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request_id.as_deref(), Some("open")); + let admission = admit_pending_requests( + pending, + &[], + 1, + 16, + 8, + 8, + 128, + |req| req.prompt_tokens.len(), + |req| req.max_tokens, + ); + assert_eq!(admission.pending.len(), 1); + assert!(admission.deferred.is_empty()); + assert!(admission.rejected.is_empty()); +} + +#[test] +fn closed_resident_work_is_absent_from_post_prune_load() { + let (closed_active_sink, closed_active_rx) = TokenSink::standalone(); + drop(closed_active_rx); + let (open_active_sink, _open_active_rx) = TokenSink::standalone(); + let (closed_prefill_sink, closed_prefill_rx) = TokenSink::standalone(); + drop(closed_prefill_rx); + let (pending_sink, _pending_rx) = TokenSink::standalone(); + let mut active = vec![ + active_request(10, "active-closed", closed_active_sink), + active_request(11, "active-open", open_active_sink), + ]; + let mut prefilling = vec![prefilling_request( + 12, + "prefill-closed", + closed_prefill_sink, + )]; + let mut pending = vec![test_request("pending-open", pending_sink)]; + let mut backend = PruneTestBackend::default(); + + assert!( + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() + ); + + assert_eq!(active.len(), 1); + assert_eq!(active[0].request_id.as_deref(), Some("active-open")); + assert!(prefilling.is_empty()); + assert_eq!( + logical_load_counts(&active, &prefilling, 0, pending.len()), + (1, 1) + ); + assert_eq!(backend.retired_active, vec![RequestId::new(10)]); + assert_eq!( + backend.dropped_prefilling, + vec![(RequestId::new(12), DropExpectation::MustBeAbsent)] + ); +} + +#[test] +fn closed_resident_frees_capacity_for_same_tick_admission() { + let (closed_sink, closed_rx) = TokenSink::standalone(); + drop(closed_rx); + let (pending_sink, _pending_rx) = TokenSink::standalone(); + let mut active = vec![active_request(20, "resident-closed", closed_sink)]; + let mut prefilling = Vec::new(); + let mut pending = vec![test_request("replacement", pending_sink)]; + let mut backend = PruneTestBackend::default(); + + assert!( + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() + ); + + let active_budget: Vec = active + .iter() + .map(|req| ActiveKvBudget { + prompt_len: req.prompt_len, + generated_count: req.generated_count, + max_tokens: req.max_tokens, + }) + .collect(); + let admission = admit_pending_requests( + pending, + &active_budget, + 1usize.saturating_sub(prefilling.len()), + 16, + 8, + 8, + 128, + |req| req.prompt_tokens.len(), + |req| req.max_tokens, + ); + + assert!(active.is_empty()); + assert_eq!(backend.retired_active, vec![RequestId::new(20)]); + assert_eq!(admission.pending.len(), 1); + assert_eq!( + admission.pending[0].request_id.as_deref(), + Some("replacement") + ); + assert!(admission.deferred.is_empty()); +} + +#[test] +fn closed_materialized_prefill_requires_existing_worker_state() { + let (closed_sink, closed_rx) = TokenSink::standalone(); + drop(closed_rx); + let mut prefilling = vec![PrefillingRequest35 { + cursor: 1, + ..prefilling_request(21, "prefill-materialized", closed_sink) + }]; + let mut active = Vec::new(); + let mut pending = Vec::new(); + let mut backend = PruneTestBackend::default(); + + assert!( + prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() + ); + + assert!(prefilling.is_empty()); + assert_eq!( + backend.dropped_prefilling, + vec![(RequestId::new(21), DropExpectation::MustExist)] + ); +} + +#[test] +fn prune_drop_failure_preserves_pending_for_terminal_fanout() { + let (closed_tx, closed_rx) = TokenSink::standalone(); + drop(closed_rx); + let (pending_tx, mut pending_rx) = TokenSink::standalone(); + let (_observer_tx, observer_rx) = TokenSink::standalone(); + let mut active = vec![active_request(22, "closed-active", closed_tx)]; + let mut prefilling = Vec::new(); + let mut pending = vec![test_request("live-pending", pending_tx)]; + let mut backend = LifecycleTestBackend::new(None, observer_rx); + backend.observer = None; + backend.fail_active_drop = true; + + let failure = + match prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) { + Ok(()) => panic!("injected prune drop should fail"), + Err(failure) => failure, + }; + assert!(active.is_empty()); + assert_eq!(failure.transient.len(), 1); + assert_eq!(pending.len(), 1); + + let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); + let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + 64, + active, + prefilling, + pending, + Vec::new(), + None, + failure, + ); + + assert!(matches!( + next_event(&mut pending_rx, "pending after prune failure"), + TokenEvent::Error { .. } + )); + assert_no_more_events(&mut pending_rx); +} + +#[test] +fn decode_eos_waits_for_drop_before_finished() { + let (token_tx, token_rx) = TokenSink::standalone(); + let mut request = active_request(30, "decode-eos", token_tx); + request.params.ignore_eos = false; + let mut active = vec![request]; + let mut backend = LifecycleTestBackend::new(Some(9), token_rx); + + assert!(dispatch_decode_tokens(&mut backend, &mut active, &[9], &[None]).is_ok()); + + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(30)]); + let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + next_event(&mut token_rx, "decode EOS"), + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + .. + } + )); + assert_no_more_events(&mut token_rx); +} + +#[test] +fn decode_length_waits_for_drop_before_token_and_finished() { + let (token_tx, token_rx) = TokenSink::standalone(); + let mut request = active_request(31, "decode-length", token_tx); + request.params.ignore_eos = true; + request.max_tokens = 2; + let mut active = vec![request]; + let mut backend = LifecycleTestBackend::new(None, token_rx); + + assert!(dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]).is_ok()); + + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(31)]); + let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + next_event(&mut token_rx, "decode length token"), + TokenEvent::Token { id: 7, .. } + )); + assert!(matches!( + next_event(&mut token_rx, "decode length finish"), + TokenEvent::Finished { + finish_reason: FinishReason::Length, + .. + } + )); + assert_no_more_events(&mut token_rx); +} + +#[test] +fn non_tp_decode_preserves_publish_before_retire_order() { + let (token_tx, token_rx) = TokenSink::standalone(); + let mut request = active_request(36, "single-order", token_tx); + request.params.ignore_eos = true; + request.max_tokens = 2; + let mut active = vec![request]; + let mut backend = LifecycleTestBackend::new(None, token_rx); + backend.active_completion_requires_drop_ack = false; + + assert!(dispatch_decode_tokens(&mut backend, &mut active, &[8], &[None]).is_ok()); + + assert!(active.is_empty()); + assert_eq!(backend.active_drops, vec![RequestId::new(36)]); + assert_eq!(backend.active_events_before_drop.len(), 2); + assert!(matches!( + &backend.active_events_before_drop[0], + TokenEvent::Token { id: 8, .. } + )); + assert!(matches!( + &backend.active_events_before_drop[1], + TokenEvent::Finished { + finish_reason: FinishReason::Length, + .. + } + )); +} + +#[test] +fn decode_completion_drop_failure_publishes_only_terminal_error() { + let (token_tx, token_rx) = TokenSink::standalone(); + let mut request = active_request(32, "decode-drop-failure", token_tx); + request.params.ignore_eos = true; + request.max_tokens = 2; + let mut active = vec![request]; + let mut backend = LifecycleTestBackend::new(None, token_rx); + backend.fail_active_drop = true; + + let failure = match dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]) { + Ok(()) => panic!("injected active drop should fail"), + Err(failure) => failure, + }; + assert!(active.is_empty()); + assert_eq!(failure.transient.len(), 1); + let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + token_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + + let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); + let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + 64, + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + None, + failure, + ); + + assert!(matches!( + next_event(&mut token_rx, "failed decode completion"), + TokenEvent::Error { .. } + )); + assert_no_more_events(&mut token_rx); +} + +#[test] +fn immediate_prefill_completion_waits_for_drop() { + let (token_tx, token_rx) = TokenSink::standalone(); + let mut request = prefilling_request(33, "prefill-length", token_tx); + request.req.max_tokens = 1; + request.req.params.ignore_eos = true; + request.step_chunk = 1; + let chunk = ScheduledChunk::from(vec![request]); + let mut active = Vec::new(); + let mut prefilling = Vec::new(); + let mut backend = LifecycleTestBackend::new(None, token_rx); + + assert!( + promote_or_requeue( + &mut backend, + &mut active, + &mut prefilling, + chunk, + &PrefillStepArtifacts::Single { + tokens: vec![11], + logprobs: vec![None], + }, + ) + .is_ok() + ); + + assert!(active.is_empty()); + assert!(prefilling.is_empty()); + assert_eq!( + backend.prefill_drops, + vec![(RequestId::new(33), DropExpectation::MustExist)] + ); + let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + next_event(&mut token_rx, "prefill length token"), + TokenEvent::Token { id: 11, .. } + )); + assert!(matches!( + next_event(&mut token_rx, "prefill length finish"), + TokenEvent::Finished { + finish_reason: FinishReason::Length, + .. + } + )); + assert_no_more_events(&mut token_rx); +} + +#[test] +fn immediate_prefill_drop_failure_publishes_only_terminal_error() { + let (token_tx, token_rx) = TokenSink::standalone(); + let (remaining_tx, mut remaining_rx) = TokenSink::standalone(); + let mut request = prefilling_request(34, "prefill-drop-failure", token_tx); + request.req.max_tokens = 1; + request.req.params.ignore_eos = true; + request.step_chunk = 1; + let mut remaining = prefilling_request(35, "remaining-scheduled", remaining_tx); + remaining.req.max_tokens = 1; + remaining.req.params.ignore_eos = true; + remaining.step_chunk = 1; + let chunk = ScheduledChunk::from(vec![request, remaining]); + let mut active = Vec::new(); + let mut prefilling = Vec::new(); + let mut backend = LifecycleTestBackend::new(None, token_rx); + backend.fail_prefill_drop = true; + + let failure = match promote_or_requeue( + &mut backend, + &mut active, + &mut prefilling, + chunk, + &PrefillStepArtifacts::Single { + tokens: vec![12, 13], + logprobs: vec![None, None], + }, + ) { + Ok(()) => panic!("injected prefill drop should fail"), + Err(failure) => failure, + }; + assert_eq!(failure.transient.len(), 2); + let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + token_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + assert!(matches!( + remaining_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + + let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); + let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + 64, + active, + prefilling, + Vec::new(), + Vec::new(), + None, + failure, + ); + + assert!(matches!( + next_event(&mut token_rx, "failed prefill completion"), + TokenEvent::Error { .. } + )); + assert_no_more_events(&mut token_rx); + assert!(matches!( + next_event(&mut remaining_rx, "remaining scheduled prefill"), + TokenEvent::Error { .. } + )); + assert_no_more_events(&mut remaining_rx); +} + +#[test] +fn terminal_shutdown_closes_drains_and_errors_every_owner_once() { + let (active_tx, active_rx) = TokenSink::standalone(); + let (prefill_tx, prefill_rx) = TokenSink::standalone(); + let (pending_tx, pending_rx) = TokenSink::standalone(); + let (deferred_tx, deferred_rx) = TokenSink::standalone(); + let (candidate_tx, candidate_rx) = TokenSink::standalone(); + let (scheduled_tx, scheduled_rx) = TokenSink::standalone(); + let (queued_tx, queued_rx) = TokenSink::standalone(); + let (after_close_tx, mut after_close_rx) = TokenSink::standalone(); + let (closed_tx, closed_rx) = TokenSink::standalone(); + drop(closed_rx); + + let active = vec![active_request(40, "active", active_tx)]; + let prefilling = vec![prefilling_request(41, "prefilling", prefill_tx)]; + let pending = vec![ + test_request("duplicate-external-id", pending_tx), + test_request("closed-sink", closed_tx), + ]; + let deferred = vec![test_request("duplicate-external-id", deferred_tx)]; + let candidate = CompletionCandidate { + request: active_request(42, "candidate", candidate_tx), + final_events: vec![TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: 1, + completion_tokens: 2, + }], + }; + let failure = FatalSchedulerError::new("injected TP replica failure") + .with_request(candidate.into_terminal()) + .with_request(test_request("scheduled", scheduled_tx)); + + let queued = test_request("queued-before-close", queued_tx); + let after_close = test_request("queued-after-close", after_close_tx); + let (submit_tx, mut submit_rx) = mpsc::unbounded_channel(); + let before_send = Arc::new(Barrier::new(2)); + let sent_before_close = Arc::new(Barrier::new(2)); + let after_receiver_close = Arc::new(Barrier::new(2)); + let sender = { + let before_send = Arc::clone(&before_send); + let sent_before_close = Arc::clone(&sent_before_close); + let after_receiver_close = Arc::clone(&after_receiver_close); + std::thread::spawn(move || { + before_send.wait(); + submit_tx + .send((queued, pegainfer_frontend::engine::KvPrefix::none())) + .expect("close-before request should be accepted"); + sent_before_close.wait(); + after_receiver_close.wait(); + submit_tx + .send((after_close, pegainfer_frontend::engine::KvPrefix::none())) + .is_err() + }) + }; + + before_send.wait(); + sent_before_close.wait(); + let (load_tx, load_rx) = watch::channel(SchedulerMetrics { + kv_used_blocks: 9, + kv_total_blocks: 64, + num_running_reqs: 9, + num_waiting_reqs: 9, + spec_decode: None, + }); + terminal_scheduler_shutdown( + &mut submit_rx, + &load_tx, + 64, + active, + prefilling, + pending, + deferred, + None, + failure, + ); + after_receiver_close.wait(); + assert!(sender.join().expect("submit race thread panicked")); + + let mut receivers = vec![ + ("active", active_rx), + ("prefilling", prefill_rx), + ("pending", pending_rx), + ("deferred", deferred_rx), + ("candidate", candidate_rx), + ("scheduled", scheduled_rx), + ("queued", queued_rx), + ]; + for (owner, rx) in &mut receivers { + match next_event(rx, owner) { + TokenEvent::Error { message, .. } => { + assert_eq!(message, "injected TP replica failure"); + } + other => panic!("{owner} received non-error terminal event: {other:?}"), + } + assert_no_more_events(rx); + } + assert!(after_close_rx.try_recv().is_err()); + + let snapshot = *load_rx.borrow(); + assert_eq!(snapshot.kv_used_blocks, 0); + assert_eq!(snapshot.kv_total_blocks, 64); + assert_eq!(snapshot.num_running_reqs, 0); + assert_eq!(snapshot.num_waiting_reqs, 0); +} + +fn collect_finished_with_timeout( + token_rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, + description: &str, +) -> (usize, FinishReason) { + let deadline = Instant::now() + Duration::from_secs(30); + let mut token_count = 0; + loop { + match token_rx.try_recv() { + Ok((_, TokenEvent::Token { .. })) => token_count += 1, + Ok((_, TokenEvent::Finished { finish_reason, .. })) => { + return (token_count, finish_reason); + } + Ok((_, TokenEvent::Error { message, .. })) => { + panic!("{description} failed: {message}") + } + Ok((_, TokenEvent::Rejected { message, .. })) => { + panic!("{description} was rejected: {message}") + } + Ok((_, _)) => {} + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { + assert!( + Instant::now() < deadline, + "timed out waiting for {description}" + ); + std::thread::sleep(Duration::from_millis(1)); + } + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { + panic!("{description} channel disconnected before Finished") + } + } + } +} + #[test] fn send_rejection_reports_kv_lifetime_request_tokens() { let (token_tx, mut token_rx) = TokenSink::standalone(); @@ -43,28 +836,44 @@ 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" - ); +fn echo_request_is_rejected_before_backend_admission() { + let (echo_tx, mut echo_rx) = TokenSink::standalone(); + let (regular_tx, mut regular_rx) = TokenSink::standalone(); + let mut echo = test_request_with_shape("unsupported-echo", echo_tx, vec![1, 2, 3], 4); + echo.echo = true; + let regular = test_request("regular", regular_tx); + let mut pending = vec![echo, regular]; + + reject_unsupported_echo(&mut pending); + + assert_eq!(pending.len(), 1); + assert_eq!(pending[0].request_id.as_deref(), Some("regular")); assert!( - matches!( - build_eager_only_plan::<&str>(true, vec![]), - Some(ExecutionPlan::Decode) - ), - "TP Phase 1 should decode only when no prefill chunk is scheduled" + !pending[0].echo, + "only requests eligible for backend admission may remain" ); + match echo_rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Rejected { + message, + prompt_tokens, + completion_tokens, + }) => { + assert_eq!(message, UNSUPPORTED_ECHO_MESSAGE); + assert_eq!(prompt_tokens, 3); + assert_eq!(completion_tokens, 0); + } + event => panic!("expected unsupported echo rejection, got {event:?}"), + } + assert!(matches!( + regular_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); } #[test] fn inflight_prefill_waits_instead_of_parking_after_last_decode_retires() { assert!( - !should_block_on_submit(true, true, true, true), + !should_block_on_submit(true, true), "an in-flight prefill must keep the scheduler off submit_rx.blocking_recv()" ); } @@ -91,67 +900,42 @@ fn tp_engine_rejects_cuda_graph_before_model_load() { #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] -fn tp2_scheduler_chunked_prefill_then_decode_smoke() { - let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") - .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); +fn tp2_scheduler_runs_forced_mixed_steps() { + let Some(model_path) = + crate::test_fixture::model_path_or_skip("tp2_scheduler_runs_forced_mixed_steps") + else { + return; + }; 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(); + start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1).expect("start TP2 scheduler"); + let (decode_tx, mut decode_rx) = TokenSink::standalone(); + let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); handle - .submit(SchedulerRequest { - trace_parent: None, - 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, - kv_transfer_params: None, - token_tx, - logprobs: 1, - echo: false, - }) - .expect("submit TP scheduler request"); + .submit(test_request_with_shape( + "mixed-active", + decode_tx, + vec![151_646], + 8, + )) + .expect("submit request that becomes active first"); + handle + .submit(test_request_with_shape( + "mixed-prefill", + prefill_tx, + vec![151_646, 9707], + 2, + )) + .expect("submit request that remains chunk-prefilling"); - 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 { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - 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"), - } - } + let (decode_tokens, decode_finish) = + collect_finished_with_timeout(&mut decode_rx, "mixed active request"); + let (prefill_tokens, prefill_finish) = + collect_finished_with_timeout(&mut prefill_rx, "mixed prefill request"); + assert_eq!(decode_tokens, 8); + assert_eq!(decode_finish, FinishReason::Length); + assert_eq!(prefill_tokens, 2); + assert_eq!(prefill_finish, FinishReason::Length); } #[test] diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 3c31a3b26..80eafc675 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1,7 +1,7 @@ //! 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. +//! Phase 2A adds one canonical eager unified command while retaining the +//! replicated linear-attention state layout from Phase 1. use std::collections::HashSet; use std::panic::AssertUnwindSafe; @@ -20,8 +20,6 @@ use anyhow::Result; use pegainfer_core::kv_pool::KvState; use pegainfer_frontend::sampler::SamplingParams; -#[cfg(test)] -use crate::batch_decode_graph::MAX_BATCH; use crate::config::TensorParallelConfig; use crate::decode_buffers::BatchDecodeBuffers35; use crate::executor::DecodePlan; @@ -45,6 +43,7 @@ const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_s 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; +const TRITON_AOT_DEVICE_TABLE_LEN: usize = 16; #[allow(dead_code)] enum TpWorkerCommand { @@ -54,28 +53,52 @@ enum TpWorkerCommand { RunPrefillChunks { chunks: Vec, sample_seed: u64, + start: Arc, resp: mpsc::Sender, }, RunDecodeStep { requests: Vec, sample_seed: u64, + start: Arc, resp: mpsc::Sender, }, RunUnifiedStep { + plan: TpUnifiedPlan, + start: Arc, resp: mpsc::Sender, }, DropRequest { request_id: RequestId, + start: Arc, + resp: mpsc::Sender, + }, + #[cfg(test)] + SnapshotState { resp: mpsc::Sender, }, + #[cfg(test)] + RemoveRequestStateForTest { + request_id: RequestId, + resp: mpsc::Sender, + }, + #[cfg(test)] + DisconnectForTest { + ready: mpsc::SyncSender<()>, + }, Shutdown, } #[derive(Debug)] enum TpWorkerReply { Ack, + DropAck { + existed: bool, + }, Prefill(PrefillResult), Decode(DecodeResult), + Unified(TpUnifiedResult), + #[cfg(test)] + Snapshot(WorkerStateSnapshot), } #[derive(Debug)] @@ -84,6 +107,58 @@ struct TpWorkerResponse { result: Result, } +/// Scheduler-owned lifecycle proof required from every TP rank during cleanup. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DropExpectation { + MustBeAbsent, + MustExist, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +enum TpCommandDecision { + #[default] + Pending, + Execute, + Cancel, +} + +#[derive(Default)] +struct TpCommandStartGate { + decision: Mutex, + changed: Condvar, +} + +impl TpCommandStartGate { + fn execute(&self) -> bool { + self.resolve(TpCommandDecision::Execute) + } + + fn cancel(&self) -> bool { + self.resolve(TpCommandDecision::Cancel) + } + + fn wait(&self) -> TpCommandDecision { + let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); + while *decision == TpCommandDecision::Pending { + decision = self + .changed + .wait(decision) + .unwrap_or_else(PoisonError::into_inner); + } + *decision + } + + fn resolve(&self, next: TpCommandDecision) -> bool { + let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); + if *decision != TpCommandDecision::Pending { + return false; + } + *decision = next; + self.changed.notify_all(); + true + } +} + #[derive(Default)] struct TpRuntimePoison { reason: Mutex>, @@ -113,7 +188,6 @@ impl TpRuntimePoison { pub struct Qwen35TpExecutor { workers: Vec, poison: Arc, - #[cfg(test)] world_size: usize, max_batch: usize, page_size: usize, @@ -188,16 +262,21 @@ impl TpDecodeStepItem { } } -impl Qwen35TpExecutor { - #[cfg(test)] - 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) - } +#[derive(Clone)] +pub(crate) struct TpUnifiedPlan { + pub(crate) prefill: Vec, + pub(crate) decode: Vec, + pub(crate) prefill_sample_seed: u64, + pub(crate) decode_sample_seed: u64, +} + +#[derive(Debug)] +pub(crate) struct TpUnifiedResult { + pub(crate) prefill: PrefillResult, + pub(crate) decode: DecodeResult, +} +impl Qwen35TpExecutor { pub fn from_runtime_with_capacity( model_path: &str, enable_cuda_graph: bool, @@ -220,6 +299,7 @@ impl Qwen35TpExecutor { max_batch: usize, max_prefill_tokens: usize, ) -> Result { + validate_cuda_ordinals(device_ordinals)?; anyhow::ensure!( device_ordinals.len() > 1, "Qwen3.5 TP executor requires at least two CUDA devices, got {}", @@ -348,7 +428,6 @@ impl Qwen35TpExecutor { Ok(Self { workers, poison, - #[cfg(test)] world_size, max_batch: min_rank_max_batch, page_size, @@ -396,7 +475,12 @@ impl Qwen35TpExecutor { )?; } drop(resp_tx); - wait_for_acks(resp_rx, self.workers.len(), "ping", &self.poison) + let responses = recv_runtime_responses(&resp_rx, self.world_size, "ping", &self.poison)?; + validate_dispatched_responses( + validate_ack_responses(responses, self.world_size, "ping"), + "ping", + &self.poison, + ) } pub fn execute_prefill(&self, plan: PrefillPlan<'_>) -> Result { @@ -427,20 +511,23 @@ impl Qwen35TpExecutor { !chunks.is_empty(), "Qwen3.5 TP prefill chunk command requires at least one chunk" ); + validate_prefill_chunks(chunks)?; 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) + let resp_rx = self.dispatch_mutating("prefill chunks", |start, resp| { + TpWorkerCommand::RunPrefillChunks { + chunks: chunks.clone(), + sample_seed, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "prefill chunks", &self.poison)?; + validate_dispatched_responses( + validate_prefill_responses(responses, self.world_size), + "prefill chunks", + &self.poison, + ) } pub fn execute_decode(&self, plan: DecodePlan<'_>) -> Result { @@ -473,38 +560,194 @@ impl Qwen35TpExecutor { !requests.is_empty(), "Qwen3.5 TP decode plan requires at least one request" ); + validate_decode_requests(requests)?; let requests = requests.to_vec(); + let resp_rx = self.dispatch_mutating("decode step", |start, resp| { + TpWorkerCommand::RunDecodeStep { + requests: requests.clone(), + sample_seed, + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "decode step", &self.poison)?; + validate_dispatched_responses( + validate_decode_responses(responses, self.world_size), + "decode step", + &self.poison, + ) + } + + pub(crate) fn execute_unified(&self, plan: &TpUnifiedPlan) -> Result { + self.poison.ensure_healthy()?; + validate_unified_plan(plan, self.max_batch)?; + let resp_rx = self.dispatch_mutating("unified step", |start, resp| { + TpWorkerCommand::RunUnifiedStep { + plan: plan.clone(), + start, + resp, + } + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "unified step", &self.poison)?; + validate_dispatched_responses( + validate_unified_responses(responses, self.world_size), + "unified step", + &self.poison, + ) + } + + pub(crate) fn poison_artifact_contract( + &self, + operation: &'static str, + err: &anyhow::Error, + ) -> anyhow::Error { + let reason = self.poison.poison(format!( + "invalid Qwen3.5 TP {operation} artifact set: {err:#}" + )); + anyhow::anyhow!(reason) + } + + pub fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { + self.poison.ensure_healthy()?; + let resp_rx = + self.dispatch_mutating("drop request", |start, resp| TpWorkerCommand::DropRequest { + request_id, + start, + resp, + })?; + let responses = + recv_runtime_responses(&resp_rx, self.world_size, "drop request", &self.poison)?; + validate_dispatched_responses( + validate_drop_responses(responses, self.world_size, expectation), + "drop request", + &self.poison, + ) + } + + #[cfg(test)] + fn snapshot_workers(&self) -> Result> { + self.poison.ensure_healthy()?; + self.snapshot_workers_unchecked_for_test() + } + + #[cfg(test)] + #[allow(dead_code)] + fn snapshot_workers_unchecked_for_test(&self) -> Result> { let (resp_tx, resp_rx) = mpsc::channel(); for worker in &self.workers { self.send_or_poison( worker, - TpWorkerCommand::RunDecodeStep { - requests: requests.clone(), - sample_seed, + TpWorkerCommand::SnapshotState { resp: resp_tx.clone(), }, )?; } drop(resp_tx); - wait_for_decode(resp_rx, self.workers.len(), &self.poison) + wait_for_worker_snapshots(&resp_rx, self.world_size, &self.poison) } - pub fn drop_request(&self, request_id: RequestId) -> Result<()> { + #[cfg(test)] + fn inject_prefill_dispatch_failure_for_test( + &self, + chunks: &[TpPrefillChunkItem], + fail_rank: usize, + ) -> Result<()> { self.poison.ensure_healthy()?; + anyhow::ensure!( + fail_rank < self.world_size, + "injected TP dispatch failure rank {fail_rank} is outside world size {}", + self.world_size + ); + validate_prefill_chunks(chunks)?; + let chunks = chunks.to_vec(); + dispatch_mutating_commands( + self.world_size, + "injected prefill chunks", + &self.poison, + |start, resp| TpWorkerCommand::RunPrefillChunks { + chunks: chunks.clone(), + sample_seed: 0, + start, + resp, + }, + |rank, command| { + if rank == fail_rank { + anyhow::bail!("injected dispatch failure at rank {rank}"); + } + self.workers[rank].send(command) + }, + )?; + Ok(()) + } + + #[cfg(test)] + fn remove_worker_request_state_for_test( + &self, + rank: usize, + request_id: RequestId, + ) -> Result { + self.poison.ensure_healthy()?; + let worker = self + .workers + .get(rank) + .ok_or_else(|| anyhow::anyhow!("test worker rank {rank} is out of range"))?; 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(), - }, - )?; + worker.send(TpWorkerCommand::RemoveRequestStateForTest { + request_id, + resp: resp_tx, + })?; + resp_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .map_err(|err| anyhow::anyhow!("test worker rank {rank} did not remove state: {err}")) + } + + #[cfg(test)] + fn disconnect_worker_receiver_for_test(&self, rank: usize) -> Result<()> { + self.poison.ensure_healthy()?; + let worker = self + .workers + .get(rank) + .ok_or_else(|| anyhow::anyhow!("test worker rank {rank} is out of range"))?; + let (ready_tx, ready_rx) = mpsc::sync_channel(1); + worker.send(TpWorkerCommand::DisconnectForTest { ready: ready_tx })?; + ready_rx + .recv_timeout(std::time::Duration::from_secs(5)) + .map_err(|err| anyhow::anyhow!("test worker rank {rank} did not disconnect: {err}"))?; + + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + let (resp_tx, _resp_rx) = mpsc::channel(); + if worker + .send(TpWorkerCommand::Ping { resp: resp_tx }) + .is_err() + { + return Ok(()); + } + anyhow::ensure!( + std::time::Instant::now() < deadline, + "test worker rank {rank} receiver remained connected" + ); + std::thread::yield_now(); } - drop(resp_tx); - wait_for_acks(resp_rx, self.workers.len(), "drop request", &self.poison) } + fn dispatch_mutating( + &self, + operation: &'static str, + build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, + ) -> Result> { + dispatch_mutating_commands( + self.world_size, + operation, + &self.poison, + build, + |rank, command| self.workers[rank].send(command), + ) + } + + #[cfg(test)] fn send_or_poison(&self, worker: &TpWorker, command: TpWorkerCommand) -> Result<()> { worker.send(command).map_err(|err| { let reason = self @@ -515,6 +758,31 @@ impl Qwen35TpExecutor { } } +fn dispatch_mutating_commands( + world_size: usize, + operation: &'static str, + poison: &TpRuntimePoison, + build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, + mut send: impl FnMut(usize, TpWorkerCommand) -> Result<()>, +) -> Result> { + let start = Arc::new(TpCommandStartGate::default()); + let (resp_tx, resp_rx) = mpsc::channel(); + for rank in 0..world_size { + let command = build(Arc::clone(&start), resp_tx.clone()); + if let Err(err) = send(rank, command) { + start.cancel(); + let reason = poison.poison(format!( + "failed to dispatch {operation} to TP worker rank {rank}: {err:#}" + )); + return Err(anyhow::anyhow!(reason)); + } + } + drop(resp_tx); + let resolved = start.execute(); + debug_assert!(resolved, "fresh TP command gate resolved more than once"); + Ok(resp_rx) +} + impl Drop for Qwen35TpExecutor { fn drop(&mut self) { for worker in &self.workers { @@ -747,6 +1015,14 @@ enum TpRequestPhase { Decoding, } +#[cfg(test)] +#[derive(Clone, Debug, PartialEq, Eq)] +struct WorkerStateSnapshot { + rank: usize, + request_count: usize, + requests: Vec<(RequestId, TpRequestPhase)>, +} + impl TpWorkerPrepared { fn new( rank: usize, @@ -884,32 +1160,75 @@ impl TpWorkerState { TpWorkerCommand::RunPrefillChunks { chunks, sample_seed, + start, resp, } => { - let result = self.execute_prefill_chunks(&chunks, sample_seed); - self.respond(resp, "prefill", result) + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_prefill_chunks(&chunks, sample_seed); + self.respond(resp, "prefill", result) + } } TpWorkerCommand::RunDecodeStep { requests, sample_seed, + start, resp, } => { - let result = self.execute_decode(&requests, sample_seed); - self.respond(resp, "decode", result) + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_decode(&requests, sample_seed); + self.respond(resp, "decode", result) + } } - TpWorkerCommand::RunUnifiedStep { resp } => { - let rank = self.rank; + TpWorkerCommand::RunUnifiedStep { plan, start, resp } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_unified(&plan); + self.respond(resp, "unified step", result) + } + } + TpWorkerCommand::DropRequest { + request_id, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let existed = self.drop_request(request_id); + self.respond(resp, "drop request", Ok(TpWorkerReply::DropAck { existed })) + } + } + #[cfg(test)] + TpWorkerCommand::SnapshotState { resp } => { + let snapshot = WorkerStateSnapshot { + rank: self.rank, + request_count: self.requests.len(), + requests: self + .requests + .iter() + .map(|state| (state.request_id, state.phase)) + .collect(), + }; self.respond( resp, - "unified step", - Err(anyhow::anyhow!( - "Qwen3.5 TP worker rank {rank} has no TP unified implementation yet" - )), + "snapshot state", + Ok(TpWorkerReply::Snapshot(snapshot)), ) } - TpWorkerCommand::DropRequest { request_id, resp } => { - self.drop_request(request_id); - self.respond(resp, "drop request", Ok(TpWorkerReply::Ack)) + #[cfg(test)] + TpWorkerCommand::RemoveRequestStateForTest { request_id, resp } => { + let _ = resp.send(self.drop_request(request_id)); + false + } + #[cfg(test)] + TpWorkerCommand::DisconnectForTest { ready } => { + let _ = ready.send(()); + break; } TpWorkerCommand::Shutdown => break, }; @@ -953,6 +1272,19 @@ impl TpWorkerState { chunks: &[TpPrefillChunkItem], sample_seed: u64, ) -> Result { + let requests = self.execute_prefill_rows(chunks, sample_seed)?; + if self.rank == 0 { + Ok(TpWorkerReply::Prefill(PrefillResult { requests })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn execute_prefill_rows( + &mut self, + chunks: &[TpPrefillChunkItem], + sample_seed: u64, + ) -> Result> { anyhow::ensure!( !chunks.is_empty(), "Qwen3.5 TP prefill chunk command requires at least one chunk" @@ -1001,13 +1333,7 @@ impl TpWorkerState { } } - if self.rank == 0 { - Ok(TpWorkerReply::Prefill(PrefillResult { - requests: primary_results, - })) - } else { - Ok(TpWorkerReply::Ack) - } + Ok(primary_results) } fn sample_final_prefill_chunk( @@ -1043,6 +1369,19 @@ impl TpWorkerState { requests: &[TpDecodeStepItem], sample_seed: u64, ) -> Result { + let requests = self.execute_decode_rows(requests, sample_seed)?; + if self.rank == 0 { + Ok(TpWorkerReply::Decode(DecodeResult { requests })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn execute_decode_rows( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result> { anyhow::ensure!( !requests.is_empty(), "Qwen3.5 TP decode command requires at least one request" @@ -1110,9 +1449,27 @@ impl TpWorkerState { } } + Ok(primary_results) + } + + fn execute_unified(&mut self, plan: &TpUnifiedPlan) -> Result { + validate_unified_worker_state(self, plan)?; + + // The command order is canonical across ranks. Sampling seeds are + // selected by the scheduler in decode-then-prefill order, independent + // of this forward order. + let prefill_requests = + self.execute_prefill_rows(&plan.prefill, plan.prefill_sample_seed)?; + let decode_requests = self.execute_decode_rows(&plan.decode, plan.decode_sample_seed)?; + if self.rank == 0 { - Ok(TpWorkerReply::Decode(DecodeResult { - requests: primary_results, + Ok(TpWorkerReply::Unified(TpUnifiedResult { + prefill: PrefillResult { + requests: prefill_requests, + }, + decode: DecodeResult { + requests: decode_requests, + }, })) } else { Ok(TpWorkerReply::Ack) @@ -1151,9 +1508,12 @@ impl TpWorkerState { .position(|state| state.request_id == request_id) } - fn drop_request(&mut self, request_id: RequestId) { + fn drop_request(&mut self, request_id: RequestId) -> bool { if let Some(idx) = self.request_index(request_id) { self.requests.swap_remove(idx); + true + } else { + false } } } @@ -1187,6 +1547,101 @@ fn validate_decode_requests(requests: &[TpDecodeStepItem]) -> Result<()> { Ok(()) } +fn validate_cuda_ordinals(device_ordinals: &[usize]) -> Result<()> { + let mut seen = HashSet::with_capacity(device_ordinals.len()); + for &ordinal in device_ordinals { + anyhow::ensure!( + ordinal < TRITON_AOT_DEVICE_TABLE_LEN, + "Qwen3.5 TP CUDA ordinal {ordinal} exceeds the Triton AOT device table bound {TRITON_AOT_DEVICE_TABLE_LEN}" + ); + anyhow::ensure!( + seen.insert(ordinal), + "Qwen3.5 TP CUDA ordinals must be distinct; ordinal {ordinal} appears more than once" + ); + } + Ok(()) +} + +fn validate_unified_plan(plan: &TpUnifiedPlan, max_batch: usize) -> Result<()> { + anyhow::ensure!( + !plan.prefill.is_empty(), + "Qwen3.5 TP unified plan requires at least one prefill chunk" + ); + anyhow::ensure!( + !plan.decode.is_empty(), + "Qwen3.5 TP unified plan requires at least one decode request" + ); + validate_prefill_chunks(&plan.prefill)?; + validate_decode_requests(&plan.decode)?; + anyhow::ensure!( + plan.prefill.len().saturating_add(plan.decode.len()) <= max_batch, + "Qwen3.5 TP unified plan has {} rows, exceeding scheduler capacity {max_batch}", + plan.prefill.len().saturating_add(plan.decode.len()) + ); + + let prefill_ids: HashSet<_> = plan.prefill.iter().map(|item| item.request_id).collect(); + for decode in &plan.decode { + anyhow::ensure!( + !prefill_ids.contains(&decode.request_id), + "Qwen3.5 TP unified plan request id {} appears in both prefill and decode", + decode.request_id.get() + ); + } + Ok(()) +} + +fn validate_unified_worker_state(state: &TpWorkerState, plan: &TpUnifiedPlan) -> Result<()> { + validate_unified_worker_layout(plan, state.max_batch, state.requests.len(), |request_id| { + state + .request_index(request_id) + .map(|idx| state.requests[idx].phase) + }) +} + +fn validate_unified_worker_layout( + plan: &TpUnifiedPlan, + max_batch: usize, + resident_count: usize, + mut phase_for: impl FnMut(RequestId) -> Option, +) -> Result<()> { + validate_unified_plan(plan, max_batch)?; + + let new_prefill_count = plan + .prefill + .iter() + .filter(|item| phase_for(item.request_id).is_none()) + .count(); + anyhow::ensure!( + resident_count.saturating_add(new_prefill_count) <= max_batch, + "Qwen3.5 TP unified plan would exceed worker capacity {}", + max_batch + ); + + for item in &plan.prefill { + if let Some(phase) = phase_for(item.request_id) { + anyhow::ensure!( + phase == TpRequestPhase::Prefilling, + "Qwen3.5 TP unified prefill request {} is already in decode state", + item.request_id.get() + ); + } + } + for item in &plan.decode { + let phase = phase_for(item.request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP unified decode request {} has no worker state", + item.request_id.get() + ) + })?; + anyhow::ensure!( + phase == TpRequestPhase::Decoding, + "Qwen3.5 TP unified request {} is not ready for decode", + item.request_id.get() + ); + } + Ok(()) +} + impl From for TpPrefillChunkItem { fn from(request: PrefillStepItem) -> Self { Self::new( @@ -1209,88 +1664,263 @@ impl From for TpDecodeStepItem { } } -#[allow(clippy::needless_pass_by_value)] -fn wait_for_acks( - responses: mpsc::Receiver, +fn recv_runtime_responses( + responses: &mpsc::Receiver, expected: usize, - op_name: &'static str, + operation: &'static str, poison: &TpRuntimePoison, -) -> Result<()> { +) -> Result> { + collect_runtime_responses(expected, operation, poison, || { + recv_runtime_response(responses, operation, poison) + }) +} + +fn collect_runtime_responses( + expected: usize, + operation: &'static str, + poison: &TpRuntimePoison, + mut recv_next: impl FnMut() -> Result, +) -> Result> { + let mut collected = Vec::with_capacity(expected); 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") - } + let response = recv_next()?; + if let Err(err) = &response.result { + // A failed rank may leave peers blocked in a collective, so response-set + // completeness is no longer recoverable or useful. + let reason = poison.poison(format!( + "rank {} failed during {operation}: {err:#}", + response.rank + )); + return Err(anyhow::anyhow!(reason)); } + collected.push(response); } - Ok(()) + Ok(collected) } -#[allow(clippy::needless_pass_by_value)] -fn wait_for_prefill( - responses: mpsc::Receiver, - expected: usize, +fn validate_dispatched_responses( + result: Result, + operation: &'static str, poison: &TpRuntimePoison, +) -> Result { + result.map_err(|err| { + let reason = poison.poison(format!( + "invalid Qwen3.5 TP {operation} response set: {err:#}" + )); + anyhow::anyhow!(reason) + }) +} + +fn validate_exact_rank_responses( + responses: Vec, + world_size: usize, + operation: &'static str, +) -> Result> { + anyhow::ensure!( + responses.len() == world_size, + "{operation} expected {world_size} responses, got {}", + responses.len() + ); + let mut seen_ranks = HashSet::with_capacity(world_size); + let mut replies = Vec::with_capacity(world_size); + for response in responses { + anyhow::ensure!( + response.rank < world_size, + "{operation} returned out-of-range rank {} for world size {world_size}", + response.rank + ); + anyhow::ensure!( + seen_ranks.insert(response.rank), + "{operation} returned duplicate rank {}", + response.rank + ); + replies.push((response.rank, response.result?)); + } + anyhow::ensure!( + (0..world_size).all(|rank| seen_ranks.contains(&rank)), + "{operation} response set did not contain every rank" + ); + replies.sort_unstable_by_key(|(rank, _)| *rank); + Ok(replies) +} + +#[cfg(test)] +fn validate_ack_responses( + responses: Vec, + world_size: usize, + operation: &'static str, +) -> Result<()> { + for (rank, reply) in validate_exact_rank_responses(responses, world_size, operation)? { + anyhow::ensure!( + matches!(reply, TpWorkerReply::Ack), + "{operation} rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ); + } + Ok(()) +} + +fn validate_drop_responses( + responses: Vec, + world_size: usize, + expectation: DropExpectation, +) -> Result<()> { + let mut existence = Vec::with_capacity(world_size); + for (rank, reply) in validate_exact_rank_responses(responses, world_size, "drop request")? { + let TpWorkerReply::DropAck { existed } = reply else { + anyhow::bail!( + "drop request rank {rank} returned {} instead of drop acknowledgement", + reply_name(&reply) + ); + }; + existence.push((rank, existed)); + } + let expected = expectation == DropExpectation::MustExist; + anyhow::ensure!( + existence.iter().all(|(_, existed)| *existed == expected), + "drop request expected {expectation:?}, got rank existence {existence:?}" + ); + Ok(()) +} + +fn validate_prefill_responses( + responses: Vec, + world_size: usize, ) -> 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") - } + let mut primary = None; + for (rank, reply) in validate_exact_rank_responses(responses, world_size, "prefill")? { + match (rank, reply) { + (0, TpWorkerReply::Prefill(result)) => primary = Some(result), + (0, reply) => anyhow::bail!( + "prefill rank 0 returned {} instead of primary prefill result", + reply_name(&reply) + ), + (_, TpWorkerReply::Ack) => {} + (rank, reply) => anyhow::bail!( + "prefill non-primary rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ), } } - result.ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP prefill returned no primary result")) + primary.ok_or_else(|| anyhow::anyhow!("prefill returned no primary result")) } -#[allow(clippy::needless_pass_by_value)] -fn wait_for_decode( - responses: mpsc::Receiver, - expected: usize, - poison: &TpRuntimePoison, +fn validate_decode_responses( + responses: Vec, + world_size: usize, ) -> Result { - let mut result = None; - for _ in 0..expected { - let response = recv_runtime_response(&responses, "decode", poison)?; + let mut primary = None; + for (rank, reply) in validate_exact_rank_responses(responses, world_size, "decode")? { + match (rank, reply) { + (0, TpWorkerReply::Decode(result)) => primary = Some(result), + (0, reply) => anyhow::bail!( + "decode rank 0 returned {} instead of primary decode result", + reply_name(&reply) + ), + (_, TpWorkerReply::Ack) => {} + (rank, reply) => anyhow::bail!( + "decode non-primary rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ), + } + } + primary.ok_or_else(|| anyhow::anyhow!("decode returned no primary result")) +} + +fn validate_unified_responses( + responses: Vec, + world_size: usize, +) -> Result { + let mut primary = None; + for (rank, reply) in validate_exact_rank_responses(responses, world_size, "unified step")? { + match (rank, reply) { + (0, TpWorkerReply::Unified(result)) => primary = Some(result), + (0, reply) => anyhow::bail!( + "unified step rank 0 returned {} instead of primary unified result", + reply_name(&reply) + ), + (_, TpWorkerReply::Ack) => {} + (rank, reply) => anyhow::bail!( + "unified step non-primary rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ), + } + } + primary.ok_or_else(|| anyhow::anyhow!("unified step returned no primary result")) +} + +fn reply_name(reply: &TpWorkerReply) -> &'static str { + match reply { + TpWorkerReply::Ack => "acknowledgement", + TpWorkerReply::DropAck { .. } => "drop acknowledgement", + TpWorkerReply::Prefill(_) => "prefill result", + TpWorkerReply::Decode(_) => "decode result", + TpWorkerReply::Unified(_) => "unified result", + #[cfg(test)] + TpWorkerReply::Snapshot(_) => "worker snapshot", + } +} + +#[cfg(test)] +fn wait_for_worker_snapshots( + responses: &mpsc::Receiver, + world_size: usize, + poison: &TpRuntimePoison, +) -> Result> { + let mut seen_ranks = HashSet::with_capacity(world_size); + let mut snapshots = Vec::with_capacity(world_size); + for _ in 0..world_size { + let response = recv_runtime_response(responses, "snapshot state", poison)?; + anyhow::ensure!( + response.rank < world_size, + "Qwen3.5 TP snapshot returned out-of-range rank {} for world size {world_size}", + response.rank + ); + anyhow::ensure!( + seen_ranks.insert(response.rank), + "Qwen3.5 TP snapshot returned duplicate rank {}", + response.rank + ); match response.result? { - TpWorkerReply::Ack => {} - TpWorkerReply::Decode(decode) => { + TpWorkerReply::Snapshot(snapshot) => { anyhow::ensure!( - response.rank == 0, - "Qwen3.5 TP decode returned a primary result from rank {}", + snapshot.rank == response.rank, + "Qwen3.5 TP snapshot payload rank {} does not match response rank {}", + snapshot.rank, response.rank ); anyhow::ensure!( - result.is_none(), - "Qwen3.5 TP decode returned multiple primary results" + snapshot.request_count == snapshot.requests.len(), + "Qwen3.5 TP rank {} snapshot count {} does not match {} request entries", + snapshot.rank, + snapshot.request_count, + snapshot.requests.len() ); - result = Some(decode); + snapshots.push(snapshot); + } + TpWorkerReply::Ack => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned acknowledgement") + } + TpWorkerReply::DropAck { .. } => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned drop acknowledgement") } TpWorkerReply::Prefill(_) => { - anyhow::bail!("Qwen3.5 TP decode unexpectedly returned prefill result") + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned prefill result") + } + TpWorkerReply::Decode(_) => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned decode result") + } + TpWorkerReply::Unified(_) => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned unified result") } } } - result.ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP decode returned no primary result")) + anyhow::ensure!( + (0..world_size).all(|rank| seen_ranks.contains(&rank)), + "Qwen3.5 TP snapshot response set did not contain every rank" + ); + snapshots.sort_unstable_by_key(|snapshot| snapshot.rank); + Ok(snapshots) } fn recv_runtime_response( @@ -1388,19 +2018,51 @@ mod tests { } #[test] - fn runtime_response_reports_any_rank_failure_immediately() { + fn runtime_response_failure_poisons_executor() { + let poison = TpRuntimePoison::default(); + let responses = vec![ + reply(0, TpWorkerReply::Ack), + TpWorkerResponse { + rank: 1, + result: Err(anyhow::anyhow!("rank 1 failed")), + }, + ]; + + let err = validate_dispatched_responses( + validate_ack_responses(responses, 2, "test"), + "test", + &poison, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("rank 1 failed")); + assert!(poison.ensure_healthy().is_err()); + } + + #[test] + fn runtime_response_collection_fails_fast_when_peer_never_responds() { let poison = TpRuntimePoison::default(); let (tx, rx) = mpsc::channel(); tx.send(TpWorkerResponse { - rank: 1, - result: Err(anyhow::anyhow!("rank 1 failed")), + rank: 0, + result: Err(anyhow::anyhow!("rank 0 failed")), }) .unwrap(); + let _keep_peer_channel_connected = tx; + let mut receive_attempts = 0; - let err = wait_for_acks(rx, 2, "test", &poison) - .unwrap_err() - .to_string(); - assert!(err.contains("rank 1 failed")); + let err = collect_runtime_responses(2, "test", &poison, || { + receive_attempts += 1; + rx.recv_timeout(std::time::Duration::from_millis(50)) + .map_err(|err| anyhow::anyhow!("waited for nonresponding rank: {err}")) + }) + .unwrap_err() + .to_string(); + + assert_eq!(receive_attempts, 1, "collector waited for the missing rank"); + assert!(err.contains("rank 0 failed")); + assert!(!err.contains("waited for nonresponding rank")); + assert!(poison.ensure_healthy().is_err()); } #[test] @@ -1416,36 +2078,56 @@ mod tests { 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); + fn reply(rank: usize, reply: TpWorkerReply) -> TpWorkerResponse { + TpWorkerResponse { + rank, + result: Ok(reply), + } } #[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 - ); - } + fn mutating_partial_dispatch_cancels_delivered_prefix_and_poisons() { + let poison = TpRuntimePoison::default(); + let (rank0_tx, rank0_rx) = mpsc::channel(); + let (rank1_tx, rank1_rx) = mpsc::channel::(); + let senders = [rank0_tx, rank1_tx]; + let err = dispatch_mutating_commands( + 2, + "test prefill", + &poison, + |start, resp| TpWorkerCommand::RunPrefillChunks { + chunks: vec![TpPrefillChunkItem::new( + RequestId::new(2), + vec![9707], + 0, + true, + )], + sample_seed: 0, + start, + resp, + }, + |rank, command| { + if rank == 1 { + anyhow::bail!("injected prefix-only dispatch failure"); + } + senders[rank] + .send(command) + .map_err(|_| anyhow::anyhow!("test receiver disconnected")) + }, + ) + .unwrap_err() + .to_string(); - #[test] - fn zero_sized_recurrent_state_keeps_requested_capacity() { - assert_eq!( - effective_recurrent_capacity(64, 0, 0, usize::MAX, usize::MAX), - 64 - ); + assert!(err.contains("injected prefix-only dispatch failure")); + let TpWorkerCommand::RunPrefillChunks { start, .. } = rank0_rx.recv().unwrap() else { + panic!("expected prefill command") + }; + assert_eq!(start.wait(), TpCommandDecision::Cancel); + assert!(matches!( + rank1_rx.try_recv(), + Err(mpsc::TryRecvError::Empty) + )); + assert!(poison.ensure_healthy().is_err()); } #[test] @@ -1509,111 +2191,330 @@ mod tests { assert!(err.contains("duplicate")); } + fn assert_workers_empty(executor: &Qwen35TpExecutor) { + let snapshots = executor + .snapshot_workers() + .expect("snapshot healthy TP workers"); + assert_snapshots_empty(&snapshots, executor.world_size()); + } + + fn assert_snapshots_empty(snapshots: &[WorkerStateSnapshot], world_size: usize) { + assert_eq!(snapshots.len(), world_size); + for (rank, snapshot) in snapshots.iter().enumerate() { + assert_eq!(snapshot.rank, rank); + assert_eq!(snapshot.request_count, 0, "rank {rank} retained requests"); + assert!( + snapshot.requests.is_empty(), + "rank {rank} retained request IDs" + ); + } + } + #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] - fn starts_tp2_workers_and_broadcasts_lifecycle_commands() { - let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") - .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + fn tp2_drop_expectations_detect_rank_lifecycle_divergence() { + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_drop_expectations_detect_rank_lifecycle_divergence", + ) else { + return; + }; 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"); + .drop_request(RequestId::new(400), DropExpectation::MustBeAbsent) + .expect("pre-materialization drop should observe all ranks absent"); + executor.ping_all().expect("absent drop preserves health"); + + let clean_id = RequestId::new(401); + executor + .execute_prefill(PrefillPlan { + requests: &[PrefillStepItem::new(clean_id, vec![151_646, 9707], 0)], + }) + .expect("materialize clean request"); + executor + .drop_request(clean_id, DropExpectation::MustExist) + .expect("materialized drop should observe all ranks present"); + assert_workers_empty(&executor); + + let divergent_id = RequestId::new(402); + executor + .execute_prefill(PrefillPlan { + requests: &[PrefillStepItem::new(divergent_id, vec![151_646, 9707], 0)], + }) + .expect("materialize divergent request"); + assert!( + executor + .remove_worker_request_state_for_test(1, divergent_id) + .expect("remove rank-1 request state") + ); + let err = executor + .drop_request(divergent_id, DropExpectation::MustExist) + .unwrap_err() + .to_string(); + assert!(err.contains("MustExist")); + assert!(executor.ping_all().is_err()); + let snapshots = executor + .snapshot_workers_unchecked_for_test() + .expect("snapshot workers after mixed drop poison"); + assert_snapshots_empty(&snapshots, executor.world_size()); } #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] - fn tp2_default_capacity_is_memory_safe() { - let model_path = std::env::var("PEGAINFER_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); + fn tp2_partial_dispatch_gate_prevents_rank_local_mutation() { + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_partial_dispatch_gate_prevents_rank_local_mutation", + ) else { + return; + }; + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + let chunk = TpPrefillChunkItem::new(RequestId::new(410), vec![151_646, 9707], 0, true); + + let err = executor + .inject_prefill_dispatch_failure_for_test(&[chunk], 1) + .unwrap_err() + .to_string(); + assert!(err.contains("injected dispatch failure at rank 1")); + assert!(executor.ping_all().is_err()); + let snapshots = executor + .snapshot_workers_unchecked_for_test() + .expect("snapshot workers after cancelled prefix dispatch"); + assert_snapshots_empty(&snapshots, executor.world_size()); } #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] - fn tp2_prefill_runs_and_returns_primary_result() { - let model_path = std::env::var("PEGAINFER_TEST_MODEL_PATH") - .unwrap_or_else(|_| "/home/data/mgj/qwen35weights".to_string()); + fn tp2_worker_receiver_disconnect_poisons_without_snapshot_claim() { + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_worker_receiver_disconnect_poisons_without_snapshot_claim", + ) else { + return; + }; 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 + executor + .disconnect_worker_receiver_for_test(1) + .expect("disconnect rank-1 worker receiver"); + + let err = executor .execute_prefill(PrefillPlan { - requests: &[request], + requests: &[PrefillStepItem::new( + RequestId::new(420), + vec![151_646, 9707], + 0, + )], }) - .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"); + .unwrap_err() + .to_string(); + assert!(err.contains("failed to dispatch prefill chunks to TP worker rank 1")); + assert!(executor.ping_all().is_err()); } #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] - fn tp2_chunked_prefill_advances_existing_request_state() { - let model_path = std::env::var("PEGAINFER_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) + fn tp2_unified_step_advances_prefill_and_decode_together() { + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_unified_step_advances_prefill_and_decode_together", + ) else { + return; + }; + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 2) .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); + let decode_id = RequestId::new(30); + let decode_prefill = executor + .execute_prefill(PrefillPlan { + requests: &[PrefillStepItem::new(decode_id, vec![151_646, 9707], 1)], + }) + .expect("materialize TP2 decode request"); + let prefill_id = RequestId::new(31); + let unified = executor + .execute_unified(&TpUnifiedPlan { + prefill: vec![TpPrefillChunkItem::new( + prefill_id, + vec![151_646, 9707], + 1, + true, + )], + decode: vec![TpDecodeStepItem::new( + decode_id, + decode_prefill.requests[0].first_token, + 1, + SamplingParams::default(), + )], + prefill_sample_seed: 102, + decode_sample_seed: 101, + }) + .expect("run TP2 unified step"); + + assert_eq!(unified.prefill.requests.len(), 1); + assert_eq!(unified.prefill.requests[0].request_id, prefill_id); + assert!(unified.prefill.requests[0].first_token_logprob.is_some()); + assert_eq!(unified.decode.requests.len(), 1); + assert_eq!(unified.decode.requests[0].request_id, decode_id); + assert!(unified.decode.requests[0].logprob.is_some()); + for snapshot in executor.snapshot_workers().expect("snapshot unified state") { + assert_eq!(snapshot.request_count, 2); + assert!( + snapshot + .requests + .iter() + .all(|(_, phase)| *phase == TpRequestPhase::Decoding) + ); + } - executor - .drop_request(request_id) - .expect("drop chunk-prefilled request"); + for request_id in [decode_id, prefill_id] { + executor + .drop_request(request_id, DropExpectation::MustExist) + .expect("drop unified request"); + } + assert_workers_empty(&executor); } #[test] #[ignore = "requires two CUDA devices and Qwen3.5 weights"] - fn tp2_decode_runs_after_prefill() { - let model_path = std::env::var("PEGAINFER_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 + fn tp2_drop_all_restores_complete_request_capacity() { + const CONFIGURED_MAX_BATCH: usize = 2; + + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_drop_all_restores_complete_request_capacity", + ) else { + return; + }; + let executor = Qwen35TpExecutor::from_runtime_with_capacity( + &model_path, + false, + &[0, 1], + CONFIGURED_MAX_BATCH, + ) + .expect("start TP2 executor"); + assert_eq!(executor.max_batch(), CONFIGURED_MAX_BATCH); + assert_workers_empty(&executor); + + let first_ids: Vec<_> = (100..100 + CONFIGURED_MAX_BATCH as u64) + .map(RequestId::new) + .collect(); + let first_requests: Vec<_> = first_ids + .iter() + .map(|&request_id| PrefillStepItem::new(request_id, vec![151_646, 9707], 0)) + .collect(); + let first_results = executor .execute_prefill(PrefillPlan { - requests: &[request], + requests: &first_requests, }) - .expect("run TP2 prefill"); - assert_eq!(prefill.requests.len(), 1); - assert_eq!(prefill.requests[0].request_id, request_id); + .expect("fill complete TP2 request capacity"); + assert_eq!(first_results.requests.len(), CONFIGURED_MAX_BATCH); + let expected_ids: HashSet<_> = first_ids.iter().copied().collect(); + for snapshot in executor + .snapshot_workers() + .expect("snapshot full TP2 request capacity") + { + assert_eq!(snapshot.request_count, CONFIGURED_MAX_BATCH); + assert_eq!( + snapshot + .requests + .iter() + .map(|(request_id, _)| *request_id) + .collect::>(), + expected_ids + ); + assert!( + snapshot + .requests + .iter() + .all(|(_, phase)| *phase == TpRequestPhase::Decoding), + "rank {} retained a non-decoding request after final prefill", + snapshot.rank + ); + } + for request_id in &first_ids { + executor + .drop_request(*request_id, DropExpectation::MustExist) + .expect("drop first-pass TP2 request"); + } + assert_workers_empty(&executor); - let decode_request = DecodeStepItem::new(request_id, prefill.requests[0].first_token, 0); + let second_ids: Vec<_> = (200..200 + CONFIGURED_MAX_BATCH as u64) + .map(RequestId::new) + .collect(); + let second_requests: Vec<_> = second_ids + .iter() + .map(|&request_id| PrefillStepItem::new(request_id, vec![151_646, 9707], 0)) + .collect(); + let second_prefill = executor + .execute_prefill(PrefillPlan { + requests: &second_requests, + }) + .expect("refill complete TP2 request capacity"); + assert_eq!(second_prefill.requests.len(), CONFIGURED_MAX_BATCH); + let decode_requests: Vec<_> = second_prefill + .requests + .iter() + .map(|result| DecodeStepItem::new(result.request_id, result.first_token, 0)) + .collect(); let decode = executor .execute_decode(DecodePlan { - requests: &[decode_request], + requests: &decode_requests, + }) + .expect("complete one decode step after TP2 capacity refill"); + assert_eq!(decode.requests.len(), CONFIGURED_MAX_BATCH); + for request_id in &second_ids { + executor + .drop_request(*request_id, DropExpectation::MustExist) + .expect("drop second-pass TP2 request"); + } + assert_workers_empty(&executor); + } + + #[test] + #[ignore = "requires two CUDA devices and Qwen3.5 weights"] + fn tp2_readmission_matches_clean_first_token_artifact() { + const REQUESTED_LOGPROBS: usize = 5; + + let Some(model_path) = crate::test_fixture::model_path_or_skip( + "tp2_readmission_matches_clean_first_token_artifact", + ) else { + return; + }; + let executor = Qwen35TpExecutor::from_runtime_with_capacity(&model_path, false, &[0, 1], 1) + .expect("start TP2 executor"); + let prompt = vec![151_646, 9707]; + + let clean_id = RequestId::new(300); + let clean_request = PrefillStepItem::new(clean_id, prompt.clone(), REQUESTED_LOGPROBS); + let clean = executor + .execute_prefill(PrefillPlan { + requests: &[clean_request], }) - .expect("run TP2 eager decode"); - assert_eq!(decode.requests.len(), 1); - assert_eq!(decode.requests[0].request_id, request_id); + .expect("run clean TP2 prefill"); + assert_eq!(clean.requests.len(), 1); + assert!(clean.requests[0].first_token_logprob.is_some()); + let clean_artifact = ( + clean.requests[0].first_token, + clean.requests[0].first_token_logprob.clone(), + ); + executor + .drop_request(clean_id, DropExpectation::MustExist) + .expect("drop clean TP2 request"); + assert_workers_empty(&executor); + let readmitted_id = RequestId::new(301); + let readmitted_request = PrefillStepItem::new(readmitted_id, prompt, REQUESTED_LOGPROBS); + let readmitted = executor + .execute_prefill(PrefillPlan { + requests: &[readmitted_request], + }) + .expect("run readmitted TP2 prefill"); + assert_eq!(readmitted.requests.len(), 1); + let readmitted_artifact = ( + readmitted.requests[0].first_token, + readmitted.requests[0].first_token_logprob.clone(), + ); + assert_eq!(readmitted_artifact, clean_artifact); executor - .drop_request(request_id) - .expect("drop decoded request"); + .drop_request(readmitted_id, DropExpectation::MustExist) + .expect("drop readmitted TP2 request"); + assert_workers_empty(&executor); } } diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index f9b0d3bdf..6526225bc 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -13,9 +13,9 @@ use std::sync::Arc; use anyhow::Result; use cudarc::driver::CudaStream; -use pegainfer_core::engine::panic_message; use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; +use pegainfer_frontend::engine::panic_message; use pegainfer_kernels::tensor::StreamOverrideGuard; use super::batch_decode_graph::BatchDecodeGraphState; @@ -138,30 +138,11 @@ impl Qwen35Model { #[cfg(test)] mod tests { - use std::path::Path; - use pegainfer_core::kv_pool::KvState; use pegainfer_core::tensor::HiddenStates; use super::*; - const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); - - fn get_model_path_or_skip() -> Option { - match std::env::var("PEGAINFER_TEST_MODEL_PATH") { - Ok(path) => Some(path), - Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { - Some(MODEL_PATH.to_string()) - } - Err(_) => { - eprintln!( - "skipping Qwen3.5 unified forward model test because {MODEL_PATH}/config.json is missing; set PEGAINFER_TEST_MODEL_PATH to run it" - ); - None - } - } - } - fn greedy_sample_batch(model: &Qwen35Model, logits: &HiddenStates, rows: usize) -> Vec { let params = vec![pegainfer_frontend::sampler::SamplingParams::default(); rows]; let params_refs: Vec<&pegainfer_frontend::sampler::SamplingParams> = @@ -177,7 +158,9 @@ mod tests { /// Verify that unified_step decode output matches batch_decode_graph standalone. #[test] fn unified_step_decode_matches_graph_decode() { - let Some(model_path) = get_model_path_or_skip() else { + let Some(model_path) = + crate::test_fixture::model_path_or_skip("unified_step_decode_matches_graph_decode") + else { return; }; let model = Qwen35Model::from_safetensors(&model_path, 0, 2).unwrap(); diff --git a/pegainfer-qwen35/tests/chunked_prefill.rs b/pegainfer-qwen35/tests/chunked_prefill.rs index 75b1629a3..dee13c8b2 100644 --- a/pegainfer-qwen35/tests/chunked_prefill.rs +++ b/pegainfer-qwen35/tests/chunked_prefill.rs @@ -18,27 +18,11 @@ use pegainfer_frontend::sampler::SamplingParams; mod common; -const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); const CHUNK_BUDGET: usize = 16; const BASELINE_PREFILL_BUDGET: usize = 1 << 20; const MAX_BATCH: usize = 2; const GENERATED_TOKENS: usize = 8; -fn model_path_or_skip() -> Option { - match std::env::var("PEGAINFER_TEST_MODEL_PATH") { - Ok(path) => Some(path), - Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { - Some(MODEL_PATH.to_string()) - } - Err(_) => { - eprintln!( - "skipping qwen35 chunked_prefill: {MODEL_PATH}/config.json is missing; set PEGAINFER_TEST_MODEL_PATH to run it" - ); - None - } - } -} - fn start_engine(model_path: &str, max_prefill_tokens: usize) -> EngineHandle { pegainfer_qwen35::start_engine( Path::new(model_path), @@ -95,7 +79,9 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec) -> (Vec, Finish #[test] fn chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv() { - let Some(model_path) = model_path_or_skip() else { + let Some(model_path) = common::model_path_or_skip( + "chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv", + ) else { return; }; let tokenizer = common::load_tokenizer(&model_path); diff --git a/pegainfer-qwen35/tests/common/mod.rs b/pegainfer-qwen35/tests/common/mod.rs index 3df2699a8..c56f1cec2 100644 --- a/pegainfer-qwen35/tests/common/mod.rs +++ b/pegainfer-qwen35/tests/common/mod.rs @@ -9,6 +9,10 @@ use vllm_text::tokenizer::HuggingFaceTokenizer; use vllm_text::tokenizer::TekkenTokenizer; use vllm_text::tokenizer::TiktokenTokenizer; +pub(crate) mod model_fixture; + +pub(crate) use model_fixture::model_path_or_skip; + #[allow(dead_code)] pub(crate) fn load_tokenizer(model_path: &str) -> DynTokenizer { try_load_tokenizer(model_path) diff --git a/pegainfer-qwen35/tests/common/model_fixture.rs b/pegainfer-qwen35/tests/common/model_fixture.rs new file mode 100644 index 000000000..bde5f640c --- /dev/null +++ b/pegainfer-qwen35/tests/common/model_fixture.rs @@ -0,0 +1,88 @@ +use std::path::Path; + +const MODEL_PATH_ENV: &str = "PEGAINFER_TEST_MODEL_PATH"; +#[allow(dead_code)] +const FRONTEND_MODEL_PATH_ENV: &str = "PEGAINFER_TEST_FRONTEND_MODEL_PATH"; + +pub(crate) fn model_path_or_skip(test_name: &str) -> Option { + fixture_path_from_env_or_skip(MODEL_PATH_ENV, test_name) +} + +#[allow(dead_code)] +pub(crate) fn frontend_model_path_or_skip( + engine_model_path: &Path, + test_name: &str, +) -> Option { + match std::env::var(FRONTEND_MODEL_PATH_ENV) { + Ok(path) => validated_fixture_path_or_skip(FRONTEND_MODEL_PATH_ENV, path, test_name), + Err(std::env::VarError::NotPresent) => { + Some(engine_model_path.to_string_lossy().into_owned()) + } + Err(std::env::VarError::NotUnicode(_)) => skip( + test_name, + &format!("{FRONTEND_MODEL_PATH_ENV} is not valid UTF-8"), + ), + } +} + +fn fixture_path_from_env_or_skip(env: &str, test_name: &str) -> Option { + match std::env::var(env) { + Ok(path) => validated_fixture_path_or_skip(env, path, test_name), + Err(std::env::VarError::NotPresent) => skip( + test_name, + &format!("{env} is not set; point it at a public Qwen3.5 model fixture"), + ), + Err(std::env::VarError::NotUnicode(_)) => { + skip(test_name, &format!("{env} is not valid UTF-8")) + } + } +} + +fn validated_fixture_path_or_skip(env: &str, path: String, test_name: &str) -> Option { + if path.trim().is_empty() { + return skip(test_name, &format!("{env} is empty")); + } + + let config_path = Path::new(&path).join("config.json"); + let raw = match std::fs::read(&config_path) { + Ok(raw) => raw, + Err(err) => { + return skip( + test_name, + &format!("cannot read {} from {env}: {err}", config_path.display()), + ); + } + }; + let config: serde_json::Value = match serde_json::from_slice(&raw) { + Ok(config) => config, + Err(err) => { + return skip( + test_name, + &format!( + "{} from {env} is not valid JSON: {err}", + config_path.display() + ), + ); + } + }; + let root_model_type = config.get("model_type").and_then(serde_json::Value::as_str); + let text_model_type = config + .pointer("/text_config/model_type") + .and_then(serde_json::Value::as_str); + if root_model_type != Some("qwen3_5") && text_model_type != Some("qwen3_5_text") { + return skip( + test_name, + &format!( + "{} from {env} is not a Qwen3.5 config", + config_path.display() + ), + ); + } + + Some(path) +} + +fn skip(test_name: &str, reason: &str) -> Option { + eprintln!("SKIP {test_name}: {reason}"); + None +} diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index d82ec7d0f..9886399bb 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -11,6 +11,7 @@ use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::SchedulerMetrics; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::TokenSink; @@ -20,8 +21,6 @@ use vllm_text::tokenizer::DynTokenizer; mod common; -const DEFAULT_MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); - const CASES: &[TestCase] = &[ TestCase { name: "tell_story", @@ -75,10 +74,6 @@ const CASES: &[TestCase] = &[ }, ]; -fn get_model_path() -> String { - std::env::var("PEGAINFER_TEST_MODEL_PATH").unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string()) -} - fn max_position_embeddings(model_path: &str) -> usize { let config_path = std::path::Path::new(model_path).join("config.json"); let config: serde_json::Value = serde_json::from_slice( @@ -226,7 +221,7 @@ fn drain_tokens(rx: &mut TokenStreamReceiver, request_id: &str) -> usize { } fn wait_for_running_requests( - load: &mut tokio::sync::watch::Receiver, + load: &mut tokio::sync::watch::Receiver, expected: u64, timeout: std::time::Duration, ) { @@ -671,7 +666,9 @@ fn run_full_scheduler_e2e( #[test] fn test_e2e_qwen35_scheduler() { - let model_path = get_model_path(); + let Some(model_path) = common::model_path_or_skip("test_e2e_qwen35_scheduler") else { + return; + }; info!("Loading Qwen3.5 model for scheduler test..."); let start = Instant::now(); @@ -696,7 +693,10 @@ fn test_e2e_qwen35_scheduler() { #[test] fn test_e2e_qwen35_shared_sm_last_decoder() { pegainfer_core::logging::init_default(); - let model_path = get_model_path(); + let Some(model_path) = common::model_path_or_skip("test_e2e_qwen35_shared_sm_last_decoder") + else { + return; + }; let tokenizer = common::load_tokenizer(&model_path); let seed_token = tokenizer .encode("Hello", false) @@ -755,7 +755,9 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { pegainfer_qwen35::Qwen35DecodeOverlap::SharedSm, ) .expect("Failed to start Qwen3.5 shared-SM scheduler"); - let mut load = handle.load_watch().expect("scheduler must expose load"); + let mut load = handle + .metrics_watch() + .expect("scheduler must expose metrics"); let mut active_rx = submit_repeated_token_request(&handle, "overlap-last-decoder", seed_token, 512, 128); @@ -820,7 +822,9 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { #[test] #[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] fn test_e2e_qwen35_scheduler_tp2() { - let model_path = get_model_path(); + let Some(model_path) = common::model_path_or_skip("test_e2e_qwen35_scheduler_tp2") else { + return; + }; info!("Loading Qwen3.5 TP2 model for scheduler test..."); let start = Instant::now(); diff --git a/pegainfer-qwen35/tests/hf_golden_gate.rs b/pegainfer-qwen35/tests/hf_golden_gate.rs index 86b269264..558866a6f 100644 --- a/pegainfer-qwen35/tests/hf_golden_gate.rs +++ b/pegainfer-qwen35/tests/hf_golden_gate.rs @@ -19,6 +19,7 @@ use std::path::PathBuf; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_qwen35::runtime::DecodePlan; use pegainfer_qwen35::runtime::DecodeStepItem; +use pegainfer_qwen35::runtime::DropExpectation; use pegainfer_qwen35::runtime::PrefillPlan; use pegainfer_qwen35::runtime::PrefillStepItem; use pegainfer_qwen35::runtime::Qwen35Executor; @@ -31,7 +32,6 @@ use sha2::Sha256; mod common; -const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); const GOLDEN_ENV: &str = "PEGAINFER_QWEN35_HF_GOLDEN"; const LONG_GOLDEN_ENV: &str = "PEGAINFER_QWEN35_HF_LONG_GOLDEN"; @@ -94,21 +94,6 @@ const BUCKET_STRADDLES: [usize; 2] = [5, 3]; const SLOT_COMPACTION_BATCH: usize = 5; const SLOT_COMPACTION_DROP_INDEX: usize = 1; -fn model_path_or_skip() -> Option { - match std::env::var("PEGAINFER_TEST_MODEL_PATH") { - Ok(path) => Some(path), - Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { - Some(MODEL_PATH.to_string()) - } - Err(_) => { - eprintln!( - "skipping qwen35 hf_golden_gate: {MODEL_PATH}/config.json is missing; set PEGAINFER_TEST_MODEL_PATH to run it" - ); - None - } - } -} - fn sha256_file(path: impl AsRef) -> Option { let bytes = std::fs::read(path).ok()?; let mut digest = Sha256::new(); @@ -575,7 +560,8 @@ fn run_tp(g: &Golden, ex: &Qwen35TpExecutor, seqs: &[usize], batched: bool) -> ( } for &id in &ids { - ex.drop_request(id).expect("TP2 drop request"); + ex.drop_request(id, DropExpectation::MustExist) + .expect("TP2 drop request"); } } else { for &seq in seqs { @@ -604,7 +590,8 @@ fn run_tp(g: &Golden, ex: &Qwen35TpExecutor, seqs: &[usize], batched: bool) -> ( &top_logprobs(dr.requests[0].logprob.as_ref()), ); } - ex.drop_request(id).expect("TP2 drop request"); + ex.drop_request(id, DropExpectation::MustExist) + .expect("TP2 drop request"); } } (stats, fingerprint) @@ -762,7 +749,7 @@ fn build_tp2_executor(model_path: &str) -> Qwen35TpExecutor { #[test] fn pega_logprobs_match_hf_golden_within_qwen35_tolerance() { - let Some(model_path) = model_path_or_skip() else { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden") else { return; }; let Some(golden) = Golden::load_for(&model_path, false) else { @@ -825,7 +812,7 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance() { #[test] fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance() { - let Some(model_path) = model_path_or_skip() else { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_long_golden") else { return; }; let Some(golden) = Golden::load_for(&model_path, true) else { @@ -850,7 +837,7 @@ fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance() { #[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 { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden_tp2") else { return; }; let Some(golden) = Golden::load_for(&model_path, false) else { @@ -879,7 +866,8 @@ fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2() { #[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 { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_long_golden_tp2") + else { return; }; let Some(golden) = Golden::load_for(&model_path, true) else { diff --git a/pegainfer-qwen35/tests/sampling_behavior.rs b/pegainfer-qwen35/tests/sampling_behavior.rs index 1dcba5a6d..4649ec186 100644 --- a/pegainfer-qwen35/tests/sampling_behavior.rs +++ b/pegainfer-qwen35/tests/sampling_behavior.rs @@ -19,24 +19,8 @@ use pegainfer_frontend::sampler::SamplingParams; mod common; -const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3.5-4B"); const GENERATED_TOKENS: usize = 32; -fn model_path_or_skip() -> Option { - match std::env::var("PEGAINFER_TEST_MODEL_PATH") { - Ok(path) => Some(path), - Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { - Some(MODEL_PATH.to_string()) - } - Err(_) => { - eprintln!( - "skipping qwen35 sampling_behavior: {MODEL_PATH}/config.json is missing; set PEGAINFER_TEST_MODEL_PATH to run it" - ); - None - } - } -} - fn params(mut params: SamplingParams) -> SamplingParams { params.ignore_eos = true; params @@ -81,7 +65,8 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingPara #[test] fn sampling_params_steer_the_qwen35_sampler() { - let Some(model_path) = model_path_or_skip() else { + let Some(model_path) = common::model_path_or_skip("sampling_params_steer_the_qwen35_sampler") + else { return; }; diff --git a/pegainfer-qwen35/tests/serving_tp2.rs b/pegainfer-qwen35/tests/serving_tp2.rs index 5d41ae7a4..2ed223e49 100644 --- a/pegainfer-qwen35/tests/serving_tp2.rs +++ b/pegainfer-qwen35/tests/serving_tp2.rs @@ -16,7 +16,6 @@ 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); @@ -39,8 +38,20 @@ impl Qwen35Tp2Server { #[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 Some(engine_model_path) = + common::model_path_or_skip("qwen35_tp2_serves_openai_completions_over_http") + else { + return Ok(()); + }; + let engine_model_path = PathBuf::from(engine_model_path); + let Some(frontend_model_path) = common::model_fixture::frontend_model_path_or_skip( + &engine_model_path, + "qwen35_tp2_serves_openai_completions_over_http", + ) else { + return Ok(()); + }; + let frontend_model_path = PathBuf::from(frontend_model_path); + let invalid_graph_model_path = engine_model_path.clone(); let server = spawn_ready_server(engine_model_path, frontend_model_path, 1).await?; let client = test_client()?; @@ -48,7 +59,11 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { 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())?; + assert_invalid_cuda_graph_tp_startup_fails( + invalid_graph_model_path + .to_str() + .context("Qwen3.5 engine fixture path is not valid UTF-8")?, + )?; server.shutdown().await } @@ -329,12 +344,3 @@ fn reserve_loopback_port() -> Result { .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("PEGAINFER_TEST_MODEL_PATH").unwrap_or_else(|_| DEFAULT_MODEL_PATH.to_string()) -} - -fn get_frontend_model_path(engine_model_path: &Path) -> String { - std::env::var("PEGAINFER_TEST_FRONTEND_MODEL_PATH") - .unwrap_or_else(|_| engine_model_path.to_string_lossy().into_owned()) -}