Conversation
mergify
Bot
added
ci/build
multi-modality
|
This pull request has merge conflicts that must be resolved before it can be |
github-project-automation
Bot
moved this from To triage
to Done
in torch.compile integration
github-project-automation
Bot
moved this from Done
to Backlog
in torch.compile integration
|
This pull request has merge conflicts that must be resolved before it can be |
zigzagcai
changed the title
9939228 to
16934c6
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
JaredforReal
left a comment
JaredforReal
left a comment
There was a problem hiding this comment.
Just curious, have u ever try launch vllm with speculative decoding? I got a little bit concern that if this feature fits for SpecDecode workload
Thanks so much for this great feature, perf gain looks great at long context prefill!
Initial Idea: we can merge #54394 first, and keep @zigzagcai 's credit at kpool version support
cc @ZJY0516
| # Every row was scored and ranked end to end by one rank, so this is | ||
| # a layout-preserving concatenation, not a top-k merge. all_gatherv | ||
| # allocates its output, so the source may alias the destination. | ||
| prefill_end = num_decode_tokens + num_prefill_tokens |
There was a problem hiding this comment.
Thanks for catching this. You are right that the original version had the same padded-tail issue.
row_shard_sizes describes the real prefill rows, while attn_metadata_narrowed.num_prefill_tokens can include CUDA Graph padding. Since all_gatherv returns sum(shard_sizes) rows, using the padded metadata count as the destination endpoint could cause a shape mismatch or overwrite padded rows.
This is fixed in the current PR head (commit c237569). The reassembly now uses:
prefill_end = num_decode_tokens + sum(shard_sizes)and only copies the gathered result into the real prefill window. The GLM k-pool path applies the same unpadded endpoint and also exchanges the expanded incomplete-pool tail (index_topk + index_kpool - 1 columns).
I also added regression coverage with padded tail rows, mixed decode/prefill rows, and the k-pool tail. The tests verify that:
- decode rows remain outside the collective;
- padded rows remain untouched;
- all real prefill rows are reconstructed correctly;
- the expanded k-pool tail is included in the exchange.
The current focused suite passes: 92 tests passed.
So this concern was valid for the earlier revision, but it is addressed in the current head following the more robust handling used in #54394.
There was a problem hiding this comment.
and we need more unit tests explict for kpool version right here
There was a problem hiding this comment.
Thanks for the invaluable feedback!
The unit test for kpool version is added in 034e20b
| # rank owns a substantial prefill slice. Keep the replicated path for | ||
| # short/medium requests; on TP4 this makes the optimized path start at 64K | ||
| # prefill rows while retaining the long-context benefit. | ||
| MIN_TP_SHARD_ROWS_PER_RANK = 16_384 |
There was a problem hiding this comment.
MIN_TP_SHARD_ROWS_PER_RANK=16K
can u give us more detail why u choose this number instead of a larger or smaller number?
There was a problem hiding this comment.
Thanks for the question. MIN_TP_SHARD_ROWS_PER_RANK is a performance amortization threshold, not a correctness requirement.
For this GLM-5.3-Flash k-pool path on H200 with TP4, we chose 16,384 rows per rank conservatively. This means that the row-sharded path is enabled only when there are at least 65,536 real scheduled prefill rows in total. Below that point, we keep the original replicated path and do not introduce the extra all_gatherv.
The threshold was raised from 1,024 after the GLM5/H200 short-to-long sweep, because the collective launch and synchronization overhead was not sufficiently amortized for smaller prefills. In the calibration sweep, the first enabled point, 64K total prefill rows, showed about 1.15% TTFT improvement and 0.94% total-token-throughput improvement. The benefit increased with longer contexts.
This is also why the 1K threshold in #54394 is not directly transferable here: that result targets a different generic DSA path, model/workload, and hardware configuration. The exact crossover is model-, kernel-, collective-, and hardware-dependent.
The current MTP-5 rerun shows the expected scaling once the sharded path is active: TTFT improvements are 3.56%, 7.78%, 15.07%, and 26.29% at 128K, 256K, 512K, and 1M respectively. We therefore kept 16K as a conservative fixed policy for this GLM5/H200 configuration rather than claiming it is a universal optimal value. The focused row-sharding test suite passes with 92 tests.
Would it be preferable to make MIN_TP_SHARD_ROWS_PER_RANK a configurable parameter, or to auto-calibrate it across different hardware configurations?
There was a problem hiding this comment.
The MIN_TP_SHARD_ROWS_PER_RANK 16K value is intentionally a conservative H200/TP4 default rather than a universal constant. The crossover depends on the MQA/top-k kernel, TP size, collective implementation, interconnect, index_topk, compression ratio, and the batch's history-length distribution.
We could keep the scope limited to this fixed safe default, since the threshold only affects performance gain and does not affect correctness. Auto-calibrating it would be useful for other hardware configurations, but could expand the scope of changes for this PR.
Or we could add runtime online auto-calibration MIN_TP_SHARD_ROWS_PER_RANK in a future PR.
vLLM uses model runnver v2 by default now. Could you also try it? |
|
This pull request has merge conflicts that must be resolved before it can be |
Hi @JaredforReal — yes, we reran the experiment with speculative decoding enabled. We used GLM-5.3-Flash on 4×H200 with TP4, Model Runner V2, The results are:
We also collected per-request SpecDecode metrics. The mean acceptance length was:
All four paired 95% confidence intervals for the acceptance-length change included zero, so we did not observe a statistically significant acceptance-rate regression. The targeted retrieval probes also had identical answer-containment rates, 48/48 on both arms. These results suggest that the optimization is compatible with the SpecDecode workload. The gain comes primarily from reducing prefill/TTFT cost; decode TPOT itself did not show a statistically significant systematic improvement, which is expected because the PR changes the prefill indexer path and does not modify the decode kernels. The focused PR tests also passed: 92 tests passed. Beyond that, I also added unit test for this feature with SpecDecode in aef70d4 |
Hi @ZJY0516 , Thanks for the invaluable feedback! The rerun used GLM-5.3-Flash on 4×H200 with TP4, Model Runner V2, explicit The V2 + MTP-5 results were:
We also collected per-request SpecDecode metrics. The mean acceptance length was:
All paired 95% confidence intervals for the acceptance-length change included zero, so we did not observe a statistically significant SpecDecode acceptance regression. The targeted retrieval probes also preserved answer containment: 48/48 on both arms. These results suggest that the row-sharding optimization is compatible with the V2 + SpecDecode workload. The improvement comes primarily from reducing prefill/TTFT cost; decode TPOT itself did not show a statistically significant systematic improvement, which is expected because this PR changes the prefill indexer path rather than the decode kernels. The focused row-sharding test suite also passed: 92 tests passed. Besides, the unittest for model runner V2 is added in commit c237569 |
Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com>
Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com>
b466281 to
aef70d4
Compare
Signed-off-by: Zheng Cai <caizheng1993@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Co-authored-by: OpenCode <noreply@opencode.ai> Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com>
ZJY0516
added
the
verified
3 participants
TL;DR
How it works / why it is faster: Instead of repeating sparse-indexer prefill MQA scoring and top-k on every TP rank, partition independent query rows into cost-balanced contiguous slices. Each rank scores its rows against the full available key history;
all_gathervreassembles the final int32 indices. This removes redundant computation in exchange for fixed-width index communication, so savings grow with context length. There is no attention approximation. Decode kernels are unchanged; smaller batches retain the existing path (activation requires at least16,384 × TPscheduled prefill query rows).Performance, baseline → optimized: main parent
5893426b88f7→ PR head16934c6e76e5; GLM-5.3-Flash on 4× H200, TP4, C1, 256 output tokens, 65,536-token batch budget, legacy runner + explicit FA3 backend. Four fresh-process A/B pairs in ABBA/BAAB order; warmups excluded. Throughput below counts input + output tokens at C1, not saturated serving throughput.Server-side TPOT at 1M: 6.631 → 6.633 ms/token. Across the full 1K–1M scan, absolute TPOT percentage-change point estimates are <0.04%; every paired 95% interval includes zero. This does not prove equivalence. The 2K TTFT point estimate is 2.84% slower (interval crosses zero); we do not claim blanket short-input non-regression.
Correctness / targeted accuracy, baseline → optimized: Across 33 distinct retrieval probes spanning 1K–1M, repeated over four processes per arm, expected-answer presence is 132/132 → 132/132; strict answer-only accuracy is 127/132 → 129/132. Exact generated-token agreement is 130/132 A/B (baseline A/A: 98/99). Both observed A/B differences contain the same correct code, with an additional explanatory sentence in baseline. These are targeted retrieval checks, not a broad quality benchmark or a claim of bitwise equivalence. The 54 focused PR tests also pass.
Full eleven-length results, paired confidence intervals, implementation scope, and reproducibility details follow.
Summary
Shard replicated sparse-indexer prefill query rows across tensor-parallel ranks. Each rank computes MQA logits and top-k for a cost-balanced contiguous slice, then one variable-size
all_gathervcall per indexer layer restores the original row layout. This removes redundant scoring; it does not approximate attention or merge independent top-k candidates. The API call is not a claim of one underlying physical NCCL operation.For GLM-5.3-Flash, the exchange includes the incomplete k-pool tail:
index_topk + index_kpool - 1 = 2,051int32 columns. KV gathering remains unconditional because subsequent query chunks can reuse that workspace.The avoided MQA scoring work grows with both query rows and the available pooled-key history, whereas the exchanged index width per query row stays fixed. Longer histories therefore offer more computation to save per byte exchanged. The cost-balanced partition also accounts for different causal-history lengths across rows. This targets prefill TTFT and prefill-dominated total-token throughput; the row partition and exchange are outside the decode branch, so a direct decode-kernel speedup is not expected. Indirect TPOT changes are measured and reported below.
Scope and activation
The diff against main is four files: the shared DSA indexer metadata/planner, the generic sparse indexer, the GLM k-pool indexer, and their focused tests. #53906 is already in the baseline; this PR does not reintroduce its model support or kernel changes.
16,384 × TPscheduled prefill query rows in the current batch.Important configuration constraint: at TP4,
max_num_batched_tokens=8192cannot activate this implementation, even with a 1M-token accumulated context. Both arms below use 65,536, so C1 inputs of 1K–32K exercise the fallback, and full 64K prefill chunks can exercise sharding. The gate is a batch-row count, not the request's context length; several shorter concurrent requests can also reach it.Exact-head A/B rescan
These measurements replace the previous historical tables. Those tables used older PR/runtime overlays and an 8,192-token batch budget; their gains cannot be attributed to the current head and gate.
5893426b88f7b3cd21101d194eb1c6f0a6f0e27b, the optimized commit's parent16934c6e76e508b22037d5d4040488d490090044Both arms load their respective clean source trees, the same model snapshot (
3f1971b7b5f7a528c9c4ef6212c8785298a8c24a), and identical compiled vLLM extensions from the official wheel for the baseline commit. Worker-level provenance checks verify the loaded vLLM paths for all four TP ranks. The historical vendor-vLLM Python overlay is not used. The PR range contains one commit.Configuration and stability protocol
VLLM_USE_V2_MODEL_RUNNER=0) and explicitFLASH_ATTN_MLA_SPARSEattention backend, chunked prefill,max_model_len=1,048,576,max_num_batched_tokens=65,536,max_num_seqs=8, GPU-memory utilization 0.90.FULL_AND_PIECEWISECUDA-graph configuration with main's default breakable-graph behavior. The effective engine configuration reportsCompilationMode.NONEand CUDA-graph capture sizes[1, 2, 4, 8, 16]; this is not a torch.compile-enabled prefill benchmark. Torch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18, Transformers 5.12.1. Same process environment/cache policy in both arms; OMP/MKL threads fixed at 1. Synchronous CUDA debugging is off for timed runs. These explicit runner/backend settings avoid failures observed in the unmodified baseline; see limitations below.ignore_eos=true. 1M means 1,000,000 input tokens, leaving room for generation within the server limit. Other K labels are powers of 1,024.1 − optimized/baseline) or higher throughput (optimized/baseline − 1). Confidence intervals resample the four restart pairs, not individual requests (20,000 bootstrap draws). Four pairs give limited statistical resolution; an interval crossing zero is not proof of equivalence.Performance
At 1M, measured TTFT improves by +25.42% and C1 total-token throughput by +32.59%. Across all lengths, TPOT improvement point estimates range from -0.03% to +0.01%. The fallback 1K–32K TTFT changes range from -2.84% to +0.43%. The intervals below and individual restart results, not just the sign of a small point estimate, determine how confidently an improvement or regression can be distinguished from run-to-run variation. These results do not prove zero overhead outside the tested configuration. Every TPOT interval includes zero: this scan does not establish a systematic TPOT change, nor does it prove exact equivalence.
Short-input caveat: the most negative fallback TTFT point is 2K: 151.39 → 155.69 ms (-2.84% improvement; paired 95% interval [-8.18, +2.24]%). This observation is retained, not discarded or declared harmless solely because the new collective is gated off. The scan does not establish blanket short-input non-regression or a causal explanation for this timing difference.
95% paired-restart bootstrap intervals for percentage improvements
The four individual adjacent-restart improvements are shown in acquisition order (AB, BA, BA, AB):
Client-side latency cross-check (same samples)
Client TTFT includes request upload/processing and transport. Server-side TPOT is the primary decode measure because clients can receive multiple generated tokens in a single SSE chunk, especially after very long prefills.
Correctness / targeted accuracy A/B
This is a targeted long-context retrieval and numerical-equivalence evaluation, not a general model-quality benchmark. Three distinct generated record probes per length place the answer at 10%, 50%, and 90% of the prompt. Each is repeated in the four independent processes per arm: 12 observations per arm per length, but only three distinct tasks. The prompt has a non-thinking assistant boundary; greedy decoding stops normally with at most 64 output tokens.
The table reports strict expected-answer accuracy, whether the expected answer is present, and exact generated-token equality for matched prompts. Answer containment alone is a weaker check: a response can contain the correct code yet fail strict answer-only formatting by adding an explanatory sentence. Baseline A/A compares adjacent independent baseline restarts, providing a control for process-to-process repeatability. HTTP success and valid token counts alone are not treated as accuracy.
Matched A/B output-token sequences agree in 130/132 comparisons; baseline A/A agrees in 98/99. The expected answer is present in 132/132 baseline and 132/132 optimized responses. Strict answer-only format passes 127/132 baseline and 129/132 optimized responses. Answer containment does not excuse extra or conflicting text; strict answer-only formatting and full token equality are reported separately. Repeated observations of the same three tasks are not independent evidence of broad model accuracy. Multimodal inputs, MTP, other DSA models, and general reasoning quality have not been evaluated by this scan.
Inspection of both A/B mismatches (the 90%-depth task at 512K and 1M) shows the same correct access code in both responses: baseline adds an explanatory sentence, while optimized returns only the code. The baseline A/A mismatch at 1M has the same formatting distinction; optimized's own restarts also differ in formatting at 512K. No different retrieved code was observed in these probes. This is an inspection of the observed mismatches, not a claim that all possible numerical differences are formatting-only.
Additional checks:
[65,536, 2,051]gathered result. Its timings are excluded from the performance table.Runtime qualification / limitations: all 704 measured performance requests complete with exactly 256 output tokens; all 264 separate retrieval requests complete. This is not an error-free-runtime claim: the common FA3 runtime emits TMA-descriptor diagnostics, which are preserved in the server logs. Two separate baseline failures were investigated before fixing the common benchmark configuration: (1) the automatically selected FlashInfer SM90 sparse MLA backend fails during 65,536-row startup/autotune; a standalone NoPE MLA probe also fails at 65,536 rows while passing 64 rows; (2) with FA3, the V2 runner fails at 256K in its slot-mapping kernel, reproduced with eager execution and CUDA_LAUNCH_BLOCKING=1. The legacy-runner/FA3 diagnostic then passes 256K, 512K, and 1M. Failed and synchronous diagnostic admissions are excluded from the performance table. No unrelated baseline source fix is overlaid. The common user environment also registers an unrelated SGLang operator at Python startup; neither pinned vLLM source tree has a callsite for that operator. These results qualify only the explicit legacy-runner/FA3 configuration, not default V2 or auto-selected FlashInfer SM90 MLA.
Reproduction
Use clean checkouts of the two pinned commits and one matched environment. Pin the compiled artifacts to the parent-main wheel (rather than allowing each editable installation to choose a different nightly):
Wheel SHA256:
93097b58f56b838b3d18928753e00b5e18624fb9084d14928e748e5fdc176dc3. The measurements used the same extracted wheel extensions in both source worktrees and checked the four workers' module paths and extension hashes.--no-depsabove assumes the explicitly pinned common runtime is already installed; it is not a complete environment bootstrap.Set
MODEL_DIRto the local model snapshot, andPYTHONto the common environment's Python. Start only one arm at a time; wait for health before running requests:Use fresh process tags
00-baseline,01-optimized,02-optimized,03-baseline,04-optimized,05-baseline,06-baseline,07-optimized. For each, run performance and then correctness with the client below; stop that owned server before starting the next. Keep the same shared compilation caches and do not include server startup in timing.Pair adjacent fresh processes and aggregate as specified above. Do not merge any diagnostic pilot or failed baseline admission into these samples. The source trees are unchanged by the benchmark; benchmark helpers are not additional PR commits.
Exact-token client and synthetic retrieval prompt generator (save as pr54951_measure.py)
Only the local model-location setting is made portable here. The tokenizer, prompt construction, seeds, token accounting, and measurement logic are the same as the measured campaign.
AI assistance and ownership
AI assistance was used for implementation and benchmark analysis. The author owns validation, review follow-up, and maintenance. This PR optimizes the already-merged GLM/DSA prefill path; it is not a duplicate model-support submission.