[Bugfix] Fix batch-invariant fp32 matmul OOR on SM89 for N=1

[vhagor]

[Bugfix] Fix batch-invariant fp32 matmul OOR on SM89 for N=1

---

Purpose

On SM86/SM89 GPUs with a ~100 KB shared-memory budget (NVIDIA L20 / L4: 101376 B/block), VLLM_BATCH_INVARIANT=1 can fail during engine warmup:

triton.runtime.errors.OutOfResources: out of resource: shared memory,
Required: 131072, Hardware limit: 101376.
Reducing block sizes or `num_stages` may help.

SM8x installs Triton matmul_persistent for aten::mm / addmm / matmul / linear. Float32 used a single hardcoded tile (128×128×32, num_stages=3). fp16 already caps BLOCK_SIZE_N from shared memory (#38670); fp32 did not.

Triton still compiles a 128-wide N-tile and masks leftover columns. On L20 + Triton 3.6.0 it then specializes on the runtime N. For N=1 that specialization adds ~32 KB of staging and the kernel needs 131072 B. Sweep N ∈ {1,2,…,4096} × K ∈ {768,2048} (plus odd K, unaligned M, bmm): only N=1 OORs.

Observed call chain on L20 warmup:

compile_or_warm_up_model
  -> dummy_pooler_run -> nn.Linear(..., 1)
  -> linear_batch_invariant -> matmul_batch_invariant
  -> matmul_persistent -> OutOfResources (131072 > 101376)

Fix: on CUDA devices with get_max_shared_memory_bytes() <= 106496, use BLOCK_SIZE_N=32, num_stages=2 only when N == 1. Every other fp32 shape keeps 128 / stages=3. BLOCK_SIZE_K stays 32. Tile choice depends on output width N (a layer property), not batch size M, so batch invariance is unchanged.

Not a duplicate of #38670 (fp16 N-tile cap). Not a duplicate of open #49131 (SM80 decode perf via M-dependent tiles; M>64 keeps the original fp32 tile, which is exactly the warmup path that OORs).

Change

In enable_batch_invariant_mode():

Device smem	fp16 BLOCK_SIZE_N	fp32 N=1 tile
> 106496 (e.g. A100)	256	128 / stages=3 (unchanged)
≤ 106496 (L20/L4)	128	32 / stages=2

Launch (matmul_persistent and bmm): fp32 uses that N=1 override only when N == 1; N > 1 stays 128 / stages=3.

Test Plan

.venv/bin/python -m pytest tests/v1/determinism/test_matmul_batch_invariant.py -v
.venv/bin/pre-commit run --files \
  vllm/model_executor/layers/batch_invariant.py \
  tests/v1/determinism/test_matmul_batch_invariant.py

New tests:

• test_fp32_n_equals_1: fp32 N=1 matches torch.mm / torch.bmm and stays batch-invariant
• test_fp32_n_gt_1_keeps_wide_tile: N>1 still matches torch.mm when the N=1 override is armed

L20 (SM 8.9, smem 101376), Triton 3.6.0:

1. Compile old vs new tile at M∈{4,128}, K=2048, N=1
2. N-sweep of the old tile
3. VLLM_BATCH_INVARIANT=1 serve BAAI/bge-reranker-base; hit /v1/models and /v1/rerank

Test Result

Kernel compile (L20, Triton 3.6.0, old 128/stages=3):

N	compiled shared	result
1	131072	OutOfResources (limit 101376)
2–15	98304	ok
≥16	81920	ok

New N=1 tile (32 / stages=2): 36864 B, launch ok (M=4 and M=128).

AI assistance

Drafting and code edits used Cursor. The submitting author reviewed every changed line and ran the L20 kernel / serve checks above.

[claude]

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.

[yewentao256]

Thanks for the work!

[yewentao256]

What is the case that N will equal to 1, in production setting?

[vhagor]

Pooling mode + classify head (num_labels=1)

[yewentao256]

LGTM, thanks for the work!

[yewentao256]

Nit: idle import now

[github-actions]

✅ @vhagor, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[mergify]

Hi @vhagor, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[vhagor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84914 for commit 2f45caaf9465.