Skip to content

[Bugfix] Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches - #56969

Merged
ZJY0516 merged 2 commits into
vllm-project:mainfrom
chaunceyjiang:glm_flashinfer_sm90
Sep 15, 2026
Merged

ZJY0516 merged 2 commits into
vllm-project:mainfrom
chaunceyjiang:glm_flashinfer_sm90

Conversation

@chaunceyjiang

@chaunceyjiang chaunceyjiang commented Sep 15, 2026

Copy link
Copy Markdown
Collaborator

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Prevent out-of-bounds access in FlashInfer SM90 sparse MLA mixed batches

Test Plan

On 4 * H20-3e

vllm serve /mnt/nvme/shared/models/ZhipuAI/GLM-5.3-Flash \
    --host 127.0.0.1 \
    --port 8000 \
    --tensor-parallel-size 4 \
    --load-format instanttensor \
    --max-model-len 8192 \
    --max-num-batched-tokens 8192 \
    --max-num-seqs 8 \
    --gpu-memory-utilization 0.85 \
    --enable-prefix-caching \
    --mamba-cache-mode align \
    --prefix-match-unit 32 \
    --additional-config '{"kda_prefill_backend":"flashkda"}' --enable-prompt-tokens-details --attention-backend FLASHINFER_MLA_SPARSE_SM90

Test Result

lm_eval run \                                                    130 ↵
  --model local-completions \
  --model_args model=/mnt/nvme/shared/models/ZhipuAI/GLM-5.3-Flash,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,tokenized_requests=False \
  --tasks gsm8k \
  --num_fewshot 5 \
  --batch_size 1 \
  --gen_kwargs temperature=0 \
  --output_path /home/chauncey/vllm/downloads/glm5.3_flash_profile/gsm8k_lm_eval \
  --log_samples \
  --seed 12345

before

(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055] WorkerProc hit an exception.
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055] Traceback (most recent call last):
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/v1/executor/multiproc_executor.py", line 1047, in _execute_worker_rpc
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     output = func(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]              ^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/v1/worker/worker_base.py", line 375, in execute_model
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self.worker.execute_model(scheduler_output)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return func(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/v1/worker/gpu_worker.py", line 1210, in execute_model
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     output = self.model_runner.execute_model(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return func(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/v1/worker/gpu/model_runner.py", line 1946, in execute_model
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     model_output = self.model(**model_inputs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]                    ^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/models/glm4_1v.py", line 2354, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     hidden_states = self.language_model.model(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/models/glm5next/nvidia/model.py", line 730, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     hidden_states, residual, post, comb = layer(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]                                           ^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/models/glm5next/nvidia/model.py", line 520, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     x = self.self_attn(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]         ^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/models/glm5next/nvidia/attention.py", line 590, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self.mla_attn(positions, hidden_states)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/mla.py", line 244, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     attn_out = self.mla_attn(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]                ^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._call_impl(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return forward_call(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 864, in forward
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     torch.ops.vllm.unified_mla_attention_with_output(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/.venv/lib/python3.12/site-packages/torch/_ops.py", line 1279, in __call__
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return self._op(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/compilation/breakable_cudagraph.py", line 104, in wrapper
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return fn(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/attention/kv_transfer_utils.py", line 40, in wrapper
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     return func(*args, **kwargs)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]            ^^^^^^^^^^^^^^^^^^^^^
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 1483, in unified_mla_attention_with_output
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     layer.forward_impl(
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 1141, in forward_impl
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     self._v_up_proj(attn_out, out=mqa_output_slice)
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]   File "/home/chauncey/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 1408, in _v_up_proj
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055]     torch.bmm(x, self.W_UV, out=out.transpose(0, 1))
(Worker_TP2 pid=1611164) ERROR 09-15 15:22:36 [multiproc_executor.py:1055] RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmStridedBatchedEx(handle, opa, opb, (int)m, (int)n, (int)k, (void*)&falpha, a, CUDA_R_16BF, (int)lda, stridea, b, CUDA_R_16BF, (int)ldb, strideb, (void*)&fbeta, c, std::is_same_v<C_Dtype, float> ? CUDA_R_32F : CUDA_R_16BF, (int)ldc, stridec, (int)num_batches, compute_type, CUBLAS_GEMM_DEFAULT_TENSOR_OP)`

this pr

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9249|±  |0.0073|
|     |       |strict-match    |     5|exact_match|↑  |0.9242|±  |0.0073|

AI assistance: Codex. Draft: human review/validation.


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

…ixed batches

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added nvidia bug Something isn't working labels Sep 15, 2026
@github-project-automation github-project-automation Bot moved this to Ready in NVIDIA Sep 15, 2026
@ZJY0516 ZJY0516 added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 15, 2026
@ZJY0516

ZJY0516 commented Sep 15, 2026

Copy link
Copy Markdown
Member

/ci run

@ZJY0516
ZJY0516 enabled auto-merge (squash) September 15, 2026 10:14
@github-actions

Copy link
Copy Markdown

❌ This PR is 10 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

@JaredforReal JaredforReal left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reproduced the bug and verified the fix works
thanks~ @chaunceyjiang

@chaunceyjiang

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #89082 for commit 395970768366.

@ZJY0516
ZJY0516 merged commit 836bb38 into vllm-project:main Sep 15, 2026
155 checks passed
@github-project-automation github-project-automation Bot moved this from Ready to Done in NVIDIA Sep 15, 2026
@chaunceyjiang
chaunceyjiang deleted the glm_flashinfer_sm90 branch September 16, 2026 06:14
drakosha added a commit to drakosha/vllm that referenced this pull request Sep 16, 2026
…and-workspace

The Indexer moved from vllm/models/glm5next/nvidia/attention.py to
vllm/models/glm5next/common/attention.py (vllm-project#55358); the prefill workspace
division by index_kpool follows it. vllm-project#56969 changed the same SM90 builder and
test file; both changes are kept.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>
keneoneth pushed a commit to keneoneth/vllm that referenced this pull request Sep 16, 2026
…ixed batches (vllm-project#56969)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working nvidia ready ONLY add when PR is ready to merge/full CI is needed

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

3 participants