Conversation
Fall back to an exact bounded-memory radix rescan when a fixed candidate buffer overflows. This prevents silent top-k corruption for narrow score distributions across the short, decode, medium, and filtered-large dispatch paths. Fixes vllm-project#51782 Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Use the first-stage histogram to detect when the threshold bin already exceeds the fixed candidate buffer. Skip the guaranteed-to-overflow collection pass and enter the existing exact radix fallback directly. On B300, five paired runs improved decode overflow cases by 6.71% and medium overflow cases by 10.90% geometrically, while modified non-overflow cases remained within 1% overall. Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
Validation overviewThis PR fixes 1. Kernel correctness and performance
2. DeepSeek-V4-Flash validation
3. GLM-5.2-NVFP4 validation
Overall, this is a real-model-triggerable silent-correctness fix with no |
|
Resolved: the build's Ninja dependency graph proves that the two modules During an independent reproducibility review, we raised a concern that the The source directory in the isolated test pod had no sibling |
|
Follow-up: the include-path concern is now resolved with the build system's recorded dependency graph.
The recorded header SHA256 values are: The two CUDA modules therefore did compile against their intended, different headers. The five-run performance comparison in my first comment remains valid. The comparison is specifically candidate versus the already-correct PR implementation, not the original incorrect vLLM |
|
Hi @yewentao256, when you have time, could you please take another look at the updated patch? I kept the correctness fallback bounded and added two targeted optimizations: vectorized filtered-row rescans and a fused decode/medium output scan. On one B300, the final code passed 435/435 correctness cases; the production-relevant 131K/k=2048 filtered-overflow case improved by 15.28%, with no measured regression on any target overflow path. The PR body now includes the full five-run methodology and results. Thank you. |
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
Thanks for the work.
Please add full reproduce command line, and e2e benchmark command line in PR description, running 3 times and adding full output log as well. Some agents will fake the number so we don't want a table only
Please also test using lm_eval to show we don't hurt acc
Vectorize exact fallback scans for long filtered rows and fuse the greater/equal output scans for decode and medium overflow paths. Preserve the two-pass filtered output path to avoid its measured long-row regression. Validated on one B300 with 435/435 correctness cases and five-seed paired benchmarks. The 131K/k=2048 filtered overflow case improves by 15.28%, while decode and medium overflow paths improve by up to 3.25% and 3.78%. Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
Adapt the long filtered-row path to an 11-bit ordered-FP16 coarse histogram and an exact 11/11/10 FP32 fallback, following the algorithmic structure of NVIDIA TensorRT-LLM production indexer Top-K radix selection: https://github.com/NVIDIA/TensorRT-LLM/blob/7a3b1bf500b742247f34d5fe4d212fd8ad284423/cpp/tensorrt_llm/kernels/indexerTopK.cu Overlay the wider histogram and CUB scan storage onto the existing 128 KiB dynamic shared-memory allocation, preserving the kernel workspace and shared-memory contract. Validated on one B300 with 435/435 correctness cases. Five-seed paired tests improve the 131K/k=2048 overflow path by 2.1357x versus the previous PR implementation, with short-path controls neutral. Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
485716d to
a01b9dd
Compare
Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
a7cc4a1 to
b98aecf
Compare
Validation and reproducerFull raw logs, Dockerfile, and exact commands: download the reproducibility bundle Conclusion: the fix passes all 435/435 overflow correctness cases. The targeted 131K/k=2048 filtered-overflow kernel improves from 0.119025 ms to 0.055705 ms, a 2.1357x median paired speedup across five runs (range 2.1299x-2.1377x), with control paths unchanged. Full-model C1/C33 tests show no measurable performance regression; the kernel is too small a fraction of E2E time to establish a statistically significant E2E speedup. GSM8K is 94.84% vs. 95.00% over 1,319 samples, within statistical error, with no material accuracy regression. unzip pr52149_raw_logs_and_reproducer.zip && cd github_raw_log_evidence
docker build -t pr52149-repro .
mkdir -p ../pr52149-work "$HOME/.cache/huggingface"
docker run --rm --gpus 'device=0' --shm-size=64g \
--ulimit memlock=-1:-1 --ulimit stack=67108864:67108864 \
-v "$PWD/../pr52149-work:/work" \
-v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
pr52149-repro -lc \
'bash /evidence/setup_repro.sh && bash /evidence/run_e2e_ab_cudagraph.sh'The PR description contains the full server, C1, C33, and lm_eval command lines. The bundle contains every formal-run stdout log; no benchmark output was omitted. |
|
Updated the PR description with the requested complete validation evidence:
Result: the overflow bug is fixed, all targeted long-vector results exactly match Complete reproducibility bundle The full per-run tables and commands are kept once in the PR description to avoid duplicating the same evidence here. @yewentao256 |
Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
LopezCastroRoberto
left a comment
LopezCastroRoberto
left a comment
There was a problem hiding this comment.
exact_topk_rescan introduces substantial duplicated traversal logic. Could you refactor the code to make the implementation easier to review and maintain? Also, tests/kernels/persistent_topk_reproducer/ should not be included in this PR; please keep the focused pytest regression and link external reproduction artifacts instead.
Signed-off-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com>
|
Thanks @LopezCastroRoberto for the review.
The final refactor passed boundary, ragged-row, vector-width, repeated-launch, and CUDA Graph correctness tests. Five paired kernel runs stayed within ±0.05% of the previous correct implementation. Compared with stock v0.27.0, C33 median TPOT was 18.6120 ms -> 17.8413 ms and output throughput was 1,076.49 -> 1,141.31 tok/s; C1 remained effectively unchanged. The refactor therefore fixes the correctness issue without hurting performance; no refactor-specific speedup is claimed. |
|
Added a real-model reproduction artifact for the affected path.
@LopezCastroRoberto |
|
Additional raw artifacts for the real GLM-5.2 overflow reproduction:
The reconstructed capture archive contains these six tensors:
Reassemble and verify: for f in glm52_real_runtime_captures.tgz.part-*.gz; do
gzip -dc "$f"
done > glm52_real_runtime_captures.tgz
echo '185c063a4d9e56c992737c64be42f8173d1f619c8f84b66765d29ac68c800f77 glm52_real_runtime_captures.tgz' | sha256sum -c -
tar -xzf glm52_real_runtime_captures.tgzThe reconstructed archive contains exactly six tensors. glm52_real_overflow_20260816_raw_logs.tgz |
|
Final validation data now supersedes the earlier intermediate performance summaries.
The PR description now uses only these final numbers. Final-head commands and raw logs: |
|
running some perf/accy evals with the proposed solution rn. will be back asap |
Can you also attach a table with the final accuracy numbers of MAIN vs PR on the MRCR eval? So we don't only know if there are collisions, but also how this affects real accuracy numbers. |
|
ok，i will give you the result asap. @LopezCastroRoberto |
|
✅ @xijiaat, CI is now available for this PR.
|
WoosukKwon
removed
the
ready
|
Thanks. I reran the end-to-end MRCR evaluation with the official
For C33, both versions completed 33/33 requests with Therefore, the final-answer accuracy is effectively unchanged in these evaluations: no regression, but also no measurable accuracy gain on these particular samples. This does not contradict the kernel-level evidence: 0.27.0 returned incorrect selected top-k values, but those errors did not alter the final answer for these samples. |
|
@xijiaat , that accuracy conclusion is exactly what I was expecting to see, because I was also not able to measure any accuracy impact. I think we need to differentiate three things here:
I did some research on this and found two interesting related discussions in SGLang: sgl-project/sglang#17747 and sgl-project/sglang#31115 (closed 3 days ago). Both are related to the overflow case. However, they weren't able to reproduce an e2e accuracy impact either, which is why the latter was eventually closed. Interestingly, the first PR explicitly mentions: "In the NSA/DSA sparse attention context, replacing a true top-k element with a numerically nearly-identical one has negligible impact on attention output." A second important aspect is that, I ran some perf analysis and, while in non-overflow cases your PR might be able to preserve performance, the added overflow path is significantly slower than the main path. You can reproduce this by using a distribution in 1.0 + torch.rand(rows, width, generator=generator) * 0.01Therefore, I tend to trust the real accuracy numbers. If there is otherwise evidence demonstrating a real impact on accuracy, then I think we should investigate that separately. But I would leave the synthetic cases aside here, because in this particular case they also have performance implications. |
|
Thanks @LopezCastroRoberto for the review. I kept this sampled-adaptive v7 experiment separate from the PR and published the complete reproducer here: https://github.com/xijiaat/persistent-topk-overflow-experiment All variants used the same production dispatcher on B300. Every cell was measured in 3 independent runs with 20 warmups, 60 samples, and 20 launches per sample.
Negative percentages mean v7 is faster. Compared with the current PR implementation, v7 is 34.66% faster over the full narrow-overflow matrix and 50.79% faster on the FilteredTopK overflow branch. Compared with MAIN, the result is mixed: normal/wide FilteredTopK is about 3% faster, while normal/wide persistent is about 4% slower. Full production-dispatch matrix: MAIN vs PR vs v7The table above contains all 504 workload cells. It reports the median kernel time across the three runs; the repository also contains every individual run value, Docker and reproduction commands, kernel sources, CSV results, and complete raw logs. v7 is correctness-clean, but it is +1,384/-39 lines versus MAIN and still has normal-path regressions, so I am keeping it as a separate optimization experiment rather than adding it to this PR. |
|
@xijiaat @LopezCastroRoberto @yewentao256 We have an exact FP32 Top-K in HPC-Ops (src, PR) that does not silently drop candidates when a coarse histogram bin overflows, and matches |
thanks！I’m still working on that！ |
@xijiaat I created #53287 to explore this problem more deeply. The goal of this PR is to produce deterministic outputs with no overhead on the non-overflow path and minimal impact when overflow occurs. It covers all three top-K variants we currently have, not only persistent top-k. My plan is to do some cleanup and break this PR down into three smaller PRs, one per kernel, to make the review process easier. Since your PR explored this issue in the Would that work for you? |
offcourse，I'm happy to work with you !! |
Fixed-size candidate buffers can silently truncate narrow score distributions and return valid but incorrect top-k indices. Detect overflow and perform an exact bounded-memory radix rescan while retaining the existing fast paths for rows that fit. Constraint: Fork-local backport containing only the persistent TopK fix; Qwen3.8-Flash-Next changes are excluded. Rejected: Open a new upstream PR | vllm-project/vllm#52149 and #55314 already address this bug class. Confidence: high Scope-risk: narrow Reversibility: clean Directive: Submit an isolated PR and merge it into the fork main branch. Tested: 13 targeted persistent TopK regression cases passed on SM120; ruff-check and clang-format passed. Not-tested: SM89 hardware and end-to-end model evaluation. Related: vllm-project/vllm#52149 Co-authored-by: ahmed xijiaat <52128022+xijiaat@users.noreply.github.com> Assisted-by: OpenAI Codex Signed-off-by: yhfgyyf <574821834@qq.com>
yewentao256
left a comment
yewentao256
left a comment
There was a problem hiding this comment.
LGTM, thanks for the work!
|
Hi @xijiaat, the pre-commit checks have failed. Please run: uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-filesThen, commit the changes and push to your branch. For future commits, |
5 participants
Purpose
This is a broad, real-model-triggerable correctness fix for
persistent_topk: candidate-buffer overflow can silently degrade attentionaccuracy, and the fix introduces no material performance regression in the
measured kernel and end-to-end workloads.
Fixes #51782.
Impact: This is not limited to synthetic logits. It reproduced end to end
on stock vLLM v0.27.0 with
nvidia/GLM-5.2-NVFP4(
GlmMoeDsaForCausalLM,index_topk=2048). A real 134,249-tokenMRCR 2-needle prompt submitted as 33 sequences dispatched
persistent_topk(rows=33,width=262144,k=2048) andsilently returned wrong top-k values in up to 27/33 rows, with up to
2,042/2,048 wrong selected values in one row and a maximum score delta of
6.5967. Replaying the same six captured runtime tensors through this PR three
times each produced zero mismatches. Longer contexts make overflow increasingly
relevant because more candidates can collide in a coarse histogram bin.
Reproduction commands and core evidence;
raw logs and tensors.
After extracting
glm52_real_overflow_20260816_raw_logs.tgz, the silent stockfailures can be located directly in the retained correctness logs:
persistent_topkuses fixed-size candidate buffers in several optimizeddispatch paths. When many scores map to the same coarse histogram bin, those
buffers can overflow. The counters keep increasing while excess candidates
are dropped, so the kernel returns valid, unique indices that are not the true
top-k values.
This change:
persistent-medium, and filtered-large paths;
overflow and be discarded;
exact
11 + 11 + 10FP32 fallback when their candidate buffer overflows;dynamic shared-memory allocation;
dynamic shared-memory launch contract;
k = 512, 1024, 2048.Algorithm provenance
The long-row radix structure is intentionally aligned with NVIDIA
TensorRT-LLM's production indexer Top-K selector. Its source explicitly uses a
half - 11 - 11 - 10 bitexact selection sequence and a CUB block scan:indexerTopK.cu, pinned source and relevant linesThis PR independently adapts that radix decomposition to vLLM's existing
filtered kernel and adds a bounded-memory exact overflow fallback; it does not
copy the TensorRT-LLM kernel. It also does not implement GVR: GVR requires
previous-step Top-K state and interface changes, which are outside this focused
correctness fix.
Exact overflow handling and optimization
Stock v0.27.0 has no exact overflow fallback: overflowing the fixed candidate
buffer can silently drop candidates and return an incorrect top-k set. This PR
detects overflow and reruns exact bounded-memory selection.
overflow frequency. If overflow remains, exact FP32 selection uses three
full-row radix passes (
11 + 11 + 10bits), preserving the dispatchVecSize(float,float2, orfloat4) and its scalar ragged tail.passes and emits values above or equal to the pivot in one fused output
scan. Tie handling remains exact.
passes followed by the shared greater-than/equal-to-pivot collectors.
for_each_score, while pivot collection is shared throughcollect_pivot_matches.histogram, and candidate buffers reuse the existing 128 KiB dynamic shared
memory allocation.
Test plan
distributions across all persistent-top-k dispatch paths.
image; no vLLM, CUDA, PyTorch, or kernel version was upgraded.
torch.topk, and validate index range and uniqueness.repetitions per case, on one NVIDIA B300 (SM 10.3, CUDA 13.0,
PyTorch 2.13.0+cu130).
capacity - 1,capacity,capacity + 1, and2 * capacityboundaries,VecSize = 1/2/4, ragged tails, 131K all-tierows, 200 repeated launches, and 200 CUDA Graph replays.
tests/kernels/test_top_k_per_row.py.preceding correct implementation with five independent paired seeds. Each
case uses 600 CUDA-event samples, 20 launches per sample, 100 warmups,
5,000 bootstrap resamples, and alternating implementation order.
Correctness result
The final PR head produced zero selected-value mismatches (
rtol=0,atol=0) across all 435/435 executions, with zero invalid indicesand zero duplicate-index rows. All targeted boundary, all-tie, ragged-row,
vector-width, repeated-launch, and CUDA Graph cases passed.
Performance result
The primary comparison is stock vLLM v0.27.0 versus the final refactored PR
head. The stock overflow path returns incorrect top-k values, so its
standalone overflow-kernel latency is not a quality-equivalent baseline.
The reviewer-requested traversal-only refactor was also measured directly
against the immediately preceding correct implementation. Across the measured
long/short control and overflow cases, five paired runs stayed within
±0.05%, so no additional refactor-specific speedup is claimed. Full-model
final-versus-stock results are reported below.
Reproduce commands
Build vLLM from this PR head, then run the checked-in overflow regression:
The standalone reproducer is intentionally kept outside this PR. After
extracting the final-refactor artifact linked below, run these commands from
its
standalone_reproducer/directory:Three-run full-model validation
Key result: the final PR head fixes all observed correctness failures with
no material performance regression in the measured workloads. Compared with
stock vLLM v0.27.0, C33 median TPOT decreased by 4.14% and output
throughput increased by 6.02%. C1 remained effectively unchanged.
The baseline was unmodified vLLM v0.27.0 and the candidate was the final
refactored PR head. They used the same vLLM image, pinned model revision,
launch flags, CUDA Graph configuration, and one isolated NVIDIA B300 each.
Every formal benchmark was run three times; warmups were excluded.
Forced-overflow correctness, three runs
For all 12 targeted long-vector case/run pairs, the PR selected-value digest
exactly matched
torch.topk(rtol=0,atol=0). All runs also had zeroinvalid indices and zero duplicate-index rows.
CUDA Graph E2E, three runs
The server command used for both binaries was:
C1 command; execute three times with seeds
2025,2026, and2027afterone excluded warmup:
C33 command; execute three times with seed
2025after one excluded requestpopulates the exact same prefix:
Median of three formal runs:
All six C33 formal runs completed 33/33 requests, generated exactly 33,000
output tokens, and reached
vllm:num_requests_running=33under 100 ms metricsampling.
GLM-5.2-NVFP4 TP2 E2E, three runs
The same final refactor was also tested end to end with the real affected model,
nvidia/GLM-5.2-NVFP4(index_topk=2048). Stock vLLM v0.27.0 and the finalPR head ran sequentially on the same isolated two B300 GPUs with identical TP2,
EP, FP8-KV, and CUDA Graph settings. Warmups were excluded.
VLLM_USE_V2_MODEL_RUNNER=1 VLLM_LOG_STATS_INTERVAL=1 \ vllm serve nvidia/GLM-5.2-NVFP4 \ --served-model-name nvidia/GLM-5.2-NVFP4 \ --host 127.0.0.1 --port 18080 --trust-remote-code \ --tensor-parallel-size 2 --enable-expert-parallel \ --moe-backend flashinfer_cutlass \ --kv-cache-dtype fp8 --gpu-memory-utilization 0.90 \ --max-model-len 262144 --max-num-batched-tokens 8192 \ --max-num-seqs 64 --enable-chunked-prefill --enable-prefix-caching \ --safetensors-load-strategy prefetch --enable-prompt-tokens-details \ --enable-request-id-headers --enable-log-requests \ --optimization-level 2 \ --compilation-config \ '{"cudagraph_mode":"FULL_AND_PIECEWISE","cudagraph_capture_sizes":[1,33]}' \ --cudagraph-metrics --seed 2025C1 was run three times with seeds
2025,2026, and2027after one excludedwarmup:
C33 was run three times with the same seed after one excluded request populated
the shared prefix:
Median of three formal runs:
All six C33 formal runs completed 33/33 requests and reached
vllm:num_requests_running=33under 100 ms metric sampling. The onlyinference-time Triton JIT occurred in the excluded C1 warmup. These results show
no GLM E2E regression and a small positive change; no large E2E speedup is
claimed.
Complete GLM E2E commands, raw stdout/server logs, detailed JSON, Prometheus
samples, and SHA256 manifests
Five full
lm_evaldatasets, three runs eachThis fresh A/B compares the final refactored PR head with rebuilt,
unmodified vLLM v0.27.0. The final fix ran first and stock ran second on the
same isolated B300 environment, model revision, server flags, and lm_eval
0.4.12 configuration. No
--limitwas used: 5,116 samples per repetition,15,348 per binary, and 30,696 total sample evaluations.
exact_match,flexible-extractexact_match,strict-matchacc,noneacc_norm,noneacc,noneprompt_level_strict_acc,noneprompt_level_loose_acc,noneinst_level_strict_acc,noneinst_level_loose_acc,noneacc,noneNo metric with an lm_eval-reported standard error shows a statistically
detectable before/after change at 95% confidence, and all observed mean deltas
are within ±0.48 percentage points. IFEval's instruction-level metrics report
no standard error upstream, so those deltas are descriptive and marked
N/Ain the complete uncertainty table.
Complete final-refactor-versus-stock accuracy artifact: all 30 stdout logs,
result JSON, sample JSONL, commands, summaries, uncertainty table, and SHA256
manifests
DeepSeek E2E and standalone reproducer artifacts
JSON, 100 ms Prometheus samples, final headers, and standalone Docker
reproducer
bundle
To run the standalone reproducer from the final-refactor artifact:
The final-refactor artifact contains every formal final-head stdout log used in
the table above. The stock bundle preserves the corresponding unmodified
v0.27.0 baseline logs. No benchmark output used by the reported medians was
omitted.