perf(qwen): shard TP4 HC prefill and coalesce recurrent checkpoints

[original-el8]

TL;DR: Add two opt-in Qwen3.8-Flash-Next prefill optimizations: distribute HyperConnection token rows across TP4 ranks, and coalesce prompt tails while exporting complete recurrent checkpoints. Together, they raise cold prefill from 3275 → 4370 tokens/s at 8K (+33.4%), 3152.5 → 3950.5 at 64K (+25.3%), and 2759 → 3330.5 at 128K (+20.7%) on four DGX Sparks. Combined decode means meet the 3% regression gate; short C8 windows limit precision. Both features default off. Coalescing requires B12X #386.

Purpose

Qwen's expanded HyperConnection (HC) state repeats substantial work on every TP rank. During eligible pure prefills, VLLM_QWEN3_8_HC_PREFILL_MODE=shard gives each rank one quarter of the token rows. HC mixing and residual updates run on those rows; all-gather supplies attention/MoE inputs, and reduce-scatter returns each block's partial output to its owning rank. PLE receives complete sequences, image deepstack residuals follow the row partition, and final sample/MTP states return in global token order.

The multimodal wrapper enters the causal model's admission path. Eligible prefills call the model forward directly, bypassing compiled wrappers. Other forwards retain full-row processing; when HC mode is enabled, target block reductions move to the decoder-layer boundary. The MTP draft layer retains its existing reductions. HC ownership uses the TP NCCL communicator; prepared RoCEnante decode collectives remain available.

VLLM_QWEN3_8_PREFILL_COALESCE=1 allows the scheduler to join a prompt tail across an internal cache boundary. GDN exports recurrent state and raw convolution history, and the companion B12X API exports PLE's normalized convolution window. Metadata tracks destination slots through cache-group remapping. Qwen MTP draft attention state is independent of target recurrent state; prefix lookup retains its existing lookahead replay requirements.

Runtime controls and compatibility

Control	Default	Behavior
VLLM_QWEN3_8_HC_PREFILL_MODE	off	off: replicated compiled path; control: replicated eager prefill with deferred reductions; shard: token ownership during eligible prefill
VLLM_QWEN3_8_PREFILL_COALESCE	0	1: one internal checkpoint per GDN/PLE request under aligned caching

HC mode requires BF16, TP4/PP1/DP1/DCP1/PCP1, no EP/EPLB, sequence-parallel MoE or DBO. Ranks agree on configuration at startup. Admission requires at least 1024 rows divisible by four, complete pure-prefill GDN metadata, and no active graph capture, full graph replay, dummy run or microbatching. Decode and mixed batches do not activate row ownership. Qualification covers GB10; other hardware is unqualified.

Coalescing requires B12X GDN prefill/decode, aligned recurrent caching and no request-boundary checkpoint mode. Install B12X #386 before enabling it. HC sharding works independently of checkpoint coalescing. The combined measured configuration enables both flags.

Checkpoint preparation

PLE preparation also primes the internal checkpoint-export executable with inactive offsets and destination slots. Coalescing can invoke this executable during graph warmup even when no request writes a checkpoint; priming it avoids a frozen-kernel-resolution failure after changing MTP depth. Prepared-call reset and close restore the live staging buffers and touched recurrent state. This changes startup preparation only.

The regression test covers coalescing enabled and disabled, inactive checkpoint sentinels, execution order, and restoration of live buffers. Both cases pass. An immutable image containing the fix passes TP4 startup, CUDA graph capture, arithmetic/tool-call checks, prefix replay, and concurrent vision smoke checks at MTP2, MTP3, and MTP4. All three depths also complete three repetitions of 8K/64K/128K cold prefill and C1/C8 short/8K decode with clean post-benchmark health checks. Each MTP depth uses a fresh Torch/vLLM cache because cached AOT graphs contain numeric prepared-plan handles. Compatible content-addressed B12X/Triton/FlashInfer caches remain reusable.

Test Result: serving performance

Status: implemented and qualified for the configuration below. Values are arithmetic mean tokens/s from two repetitions, measured on the same immutable image with runtime flags selecting each arm. The reference is that image with both flags off. All cold-prefill samples have zero cached tokens.

Cold prompt	Flags off	Eager control	HC only	Coalescing only	Combined	Combined change
8K	3275.0	3263.0	3965.5	3573.0	4370.0	+33.44%
64K	3152.5	3152.0	3911.0	3196.5	3950.5	+25.31%
128K	2759.0	2755.5	3321.0	2769.5	3330.5	+20.71%

The eager control differs from the compiled reference by less than 0.4% at every prefill length. HC ownership accounts for most of the improvement; coalescing adds its largest benefit at 8K.

Decode load	Flags off	HC only	Coalescing only	Combined	Combined change
C1, short	70.25	68.31	68.44	70.91	+0.93%
C1, 8K	62.46	62.28	61.52	64.51	+3.28%
C8, short	252.02	246.48	242.91	252.51	+0.19%
C8, 8K	183.21	179.65	197.81	204.07	+11.39%

Combined and HC-only means meet the declared 3% decode regression gate. Coalescing alone fails short-context C8 at -3.61%. Decode uses continuous requests with 256 output tokens, a three-second warmup and a 20-second observation window. C8 reaches eight active requests but averages about 7.4–7.9 during request turnover; some samples are underfilled or capacity-limited. Prompt nonces and MTP acceptance vary between runs. These short sequential measurements do not establish statistical equivalence or an isolated decode-kernel speedup. All measured arms completed without request errors, warmup timeouts, OOMs or container restarts.

Individual baseline and combined samples, tokens/s

Cell	Flags off, runs 1 / 2	Combined, runs 1 / 2
Prefill 8K	3284 / 3266	4383 / 4357
Prefill 64K	3159 / 3146	3950 / 3951
Prefill 128K	2764 / 2754	3337 / 3324
Decode C1, short	70.2115 / 70.2895	70.6982 / 71.1144
Decode C1, 8K	63.0808 / 61.8454	64.0802 / 64.9470
Decode C8, short	255.6677 / 248.3787	253.1403 / 251.8782
Decode C8, 8K	172.7899 / 193.6240	202.5287 / 205.6027

Conditions and source identity

Four DGX Sparks, NVIDIA GB10, TP4/PP1/DP1/DCP1/PCP1, Qwen3.8-Flash-Next NVFP4 with full resident PLE, InstantTensor loading, MTP3, max batched tokens 8192, max sequences 16, max model length 262144, BF16 KV and 28 GiB KV per rank. Hybrid alignment yields an effective 752-token block size. Prefix caching, chunked prefill and async scheduling are enabled; recurrent caching uses align with default retention interval zero. Decode uses full/piecewise graphs at 4/8/16/32/64 rows. RoCEnante thresholds are 2 MiB all-reduce and 16 MiB all-gather over the dual-HCA fabric.

Measured vLLM: b0cf3b82341433caa2b350b1b71fc98b11c8ed92, based on JJ 59fbf050084aefe2a2dd3b19b271a7eb3527c6ae. Measured B12X: e9653dc1eae2b7b19357c51a420bd597fe3012ed; the companion PR has the identical full source tree aed60ccc03821db784c875b8ea124acf7e0055f1 after cherry-picking onto the merged RoCEnante fixes. All ranks used image ID sha256:e2140e8359fb185a5f06e26dbc18d5f4f9ba4b3adb694fd39c1ec9948ecbb027.

Test Plan and correctness results

Source checks ran in the candidate image with CUDA_VISIBLE_DEVICES= and VLLM_TARGET_DEVICE=cpu:

.venv/bin/python -m pytest -q \
  tests/models/test_qwen3_8_flash_next_config.py \
  tests/models/test_qwen3_8_flash_next_model.py \
  tests/models/test_qwen3_8_hc_prefill.py \
  tests/v1/attention/test_ple_metadata_builder.py \
  tests/v1/core/test_mamba_align_chunk_split.py

101 unique passes, 22 GPU-only skips. Three metadata cases initially lacked the Hugging Face config fixture; rerunning the complete metadata file with the existing Qwen/Qwen3.5-0.8B config cache gave 8 passed / 2 skipped. No source change was needed. Commit pre-commit hooks pass.

The companion B12X native PLE selection passes 4 GPU tests, covering mixed requests, short history offsets, graph replay without allocation, and high destination slots beyond signed 32-bit element offsets. Serving logs confirm the expected HC activation on all four ranks: 96 reduce-scatters and 99 gathers per eligible target forward.

Combined functional checks pass arithmetic, tool use, image color, cold/warm prefix reuse and concurrent C8 replay. A 6598-token vision prompt exercises nonzero deepstack residuals during long prefill and returns the expected result on cold and warm requests. Two fresh 8192-token text prompts have zero cold cache hits, reuse 6768 tokens when warm, and reproduce their eight-token continuations.

Teacher-forced evaluation uses fixed 8192-token prose and code corpora, identical token IDs per arm, temperature zero, prompt logprobs and top-five probabilities. Gates are mean NLL increase ≤0.02 nats and top-1 agreement ≥99% where the baseline top-two margin is ≥0.2 nats:

Combined vs flags off	Prose	Code
Strong-margin top-1 agreement	99.878%	99.976%
Mean NLL increase	0.000689 nats	0.000229 nats

Both corpora pass; the separate HC and coalescing arms also pass. These checks do not establish general model quality, broad hardware support or long-duration reliability. General benchmark accuracy evaluation remains unrun.

Four-rank integer gather/reduction oracles pass. BF16 all-reduce and reduce-scatter each have about 0.26% relative RMS error against an FP64 sum, while their two-block synthetic HC outputs differ by about 0.44%. The original 0.4% cross-path threshold fails; the fixed-prompt model checks above provide separate numerical evidence and do not imply bitwise equivalence.

Serving benchmark invocation and method

The local benchmark-qwen.py adapter runs llm_decode_bench.py and fixes enable_thinking=false for generation and tokenization. Adapter SHA-256: dfdde00f72923a6aac8bf4b8111cef7d4b2ebc1bf303a67e71cf52ec2154dfe5; harness SHA-256: 053989edff8c9c93e2b96e61342b2ffbd9851e03deba17e6d3fc96fcd6694c1e. These local harness files are not included in this source PR.

.venv/bin/python benchmark-qwen.py \
  --host SERVER --port 8000 --model Qwen3.8-Flash-Next \
  --concurrency 1,8 --contexts 0,8k --prefill-contexts 8k,64k,128k \
  --standalone-prefill --prefill-duration 1 --max-tokens 256 \
  --temperature 0 --duration 20 --decode-warmup-seconds 3 \
  --cell-warmup-timeout-seconds 180 --token-targeting exact \
  --no-calibration-cache --prefill-metric auto --display-mode plain \
  --no-hw-monitor --no-resume --output RESULT.json

Prefill rates use server-measured prefill duration, with exact nominal lengths plus two chat-template tokens and explicit zero cached tokens. The baseline and combined sample table retains both repetitions. Raw responses, metrics, rank logs, source/image receipts and fixed protocols are retained in the local qualification bundle qwen38-hc-20260916; this PR reports their results without publishing the full local bundle.

Related work and review status

Open-PR checks found no duplicate Qwen TP4 HC ownership plus GDN/PLE checkpoint integration. #700 and #699 implement analogous GLM mHC ownership and KDA continuation checkpoints. This PR handles Qwen's HC layout, GDN/PLE history, multimodal residuals and MTP contract. Upstream #53909 adds standalone Qwen kernels, not this distributed ownership or scheduler integration.

AI assistance was used for implementation, testing and PR preparation. This is a draft against the maintained JJ fork. Human line-by-line review and independent validation remain pending before marking it ready to merge.

Summary by CodeRabbit

• New Features

  • Added optional row-sharded prefill execution for Qwen3.8-Flash-Next, with efficient tensor gathering and reduction.
  • Added configuration options for prefill modes and coalescing.
  • Added prefill checkpoint support for recurrent attention and PLE layers.
  • Improved Qwen3.8-Flash-Next MTP scheduling compatibility.

• Tests

  • Added coverage for prefill eligibility, row ownership, collective operations, checkpoint tracking, and staging-buffer restoration.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: e17a2762-1e7c-4c89-a4bb-5205f6da9a5d

📥 Commits

Reviewing files that changed from the base of the PR and between b0cf3b8 and d320ffd.

📒 Files selected for processing (2)

• tests/models/test_qwen3_8_flash_next_model.py
• vllm/models/qwen3_8_flash_next/ple_layer.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

The pull request adds HC prefill support for Qwen3.8 Flash Next, including TP row ownership, deferred reductions, eligibility checks, and environment configuration. It also adds PLE and GDN prefill checkpoint metadata, checkpoint storage, scheduler integration, and tests.

Changes

Qwen3.8 Flash Next prefill

Layer / File(s)	Summary
HC prefill ownership and execution vllm/envs.py, vllm/models/qwen3_8_flash_next/hc_prefill.py, vllm/models/qwen3_8_flash_next/model.py, vllm/model_executor/models/qwen3_next.py, tests/models/*	Adds HC prefill configuration, eligibility checks, TP row ownership, deferred reductions, output gathering, and related tests.
Prefill checkpoint pipeline vllm/models/qwen3_8_flash_next/ple_attn.py, vllm/models/qwen3_8_flash_next/ple_layer.py, vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py, tests/v1/attention/test_ple_metadata_builder.py, tests/models/test_qwen3_8_flash_next_model.py	Adds checkpoint columns, offsets, and slots to PLE metadata. PLE and GDN paths export checkpoint state under supported cache and backend settings.
MTP checkpoint scheduling vllm/v1/core/sched/scheduler.py	Recognizes Qwen3_8FlashNextMTP as an architecture with independent draft state for Mamba checkpoint handling.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant CausalLM
  participant hc_prefill
  participant QwenModel
  participant DecoderLayers
  participant TPGroup
  CausalLM->>hc_prefill: evaluate prefill eligibility
  hc_prefill->>QwenModel: create row ownership
  QwenModel->>DecoderLayers: run eager prefill with owner
  DecoderLayers->>TPGroup: reduce-scatter deferred results
  QwenModel->>TPGroup: gather final rows

Loading

Suggested reviewers: lukealonso

Merge Risk: ⚪ Minimal · up to d320f

No actionable current-head risk remains from the reviewed changes.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.56% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 72 functions across 11 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely identifies both primary changes: TP4 HC prefill sharding and recurrent checkpoint coalescing for Qwen.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/qwen-hc-prefill-20260916

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}