fix: correct Vulkan Linear output for batches smaller than 4 (garbage decode)

[maci0]

Problem

Token generation produces garbage after the prompt. The model loads, the prompt
is processed, but every generated token is wrong (e.g. "the the the ...").

Root cause: the Vulkan matmul kernel used by _wrap_linear tiles 4 rows at a
time and returns numerically incorrect results when the row dimension M < 4.
Decode runs one token at a time (M = 1), so every decode step computes wrong
logits while prefill (M = prompt_len >= 4) is fine — hence correct-looking
prefill but garbage generation.

Verified directly against torch.nn.functional.linear on a discrete AMD GPU
(RADV, gfx1100):

M= 1 rel_err=0.90  BAD
M= 2 rel_err=0.94  BAD
M= 3 rel_err=0.97  BAD
M= 4 rel_err=0.00  OK
M>=4 ...           OK

Fix

Pad the row dimension up to 4 in _wrap_linear, run the kernel, then slice the
result back. Single extra (duplicated) row for M < 4; no change for M >= 4.
After this, generation is coherent ("The capital of France is ... Paris ...").

This is a workaround at the dispatch layer; the underlying matmul shader's
small-M tiling could be fixed to drop it later (the added kernel test guards
the M >= 4 path and documents the expectation).

Tests

tests/python/test_linear_dispatch.py (skipped without a Vulkan device):

• wrapped Linear matches plain nn.Linear for M in 1,2,3,4,8,17, with and
  without bias
• the raw kernel stays exact for M >= 4

[gemini-code-assist]

Code Review

This pull request introduces a workaround for a Vulkan matmul kernel bug that produces incorrect results for small batches (M < 4) by padding the input to 4 rows and slicing the output. It also adds unit tests to verify correctness. However, the review identified a critical performance issue where casting the weight to float on every forward pass bypasses the weight caching mechanism, causing massive overhead by re-uploading weights to the GPU. Additionally, an empty batch (M = 0) will cause a crash. A code suggestion is provided to resolve both issues. Confidence Score: 2/5

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Critical Caching and Empty Batch Issues

1. Weight Caching Bypass (Performance Disaster): Passing weight.float() to vulkan_ops.linear creates a temporary float32 tensor with a new untyped storage on every single forward pass. Because _get_or_upload_weight uses id(weight.untyped_storage()) as the cache key, this temporary tensor results in a cache miss every single step. Consequently, the entire weight matrix is re-uploaded to the GPU on every single token generation step, completely defeating the weight caching mechanism and causing massive overhead. Passing the original weight directly avoids this, as vulkan_ops.linear and _vulkan_matvec are already designed to handle the original weight and cache it correctly.

2. Empty Batch Crash: If m == 0 (which can happen in vLLM during empty iterations or warmup phases), xf[-1:] will raise an IndexError: index -1 is out of bounds. Handling m == 0 explicitly by returning an empty tensor avoids this crash and skips unnecessary Vulkan kernel launches.

Suggested change

	if m < 4:
	xf = torch.cat([xf, xf[-1:].expand(4 - m, -1)], dim=0)
	result = vulkan_ops.linear(xf, weight.float(), bias)[:m]
	else:
	result = vulkan_ops.linear(xf, weight.float(), bias)
	if m == 0:
	result = xf.new_empty(0, weight.shape[0])
	elif m < 4:
	xf = torch.cat([xf, xf[-1:].expand(4 - m, -1)], dim=0)
	result = vulkan_ops.linear(xf, weight, bias)[:m]
	else:
	result = vulkan_ops.linear(xf, weight, bias)

References

1. Focus on correctness, safety, and maintainability — in that order. Ensure correctness and prevent crashes (safety) under empty inputs. ^(link)

[maci0]

Good catches, both fixed in the latest push:

• Empty batch (M=0): now returns an empty result before the padding path, so no more IndexError on xf[-1:] during warmup/empty iterations.
• Weight caching: switched from weight.float() to passing the original weight. Confirmed vulkan_ops.linear handles bf16/fp16/fp32 weights directly and caches by storage id, so the float cast was the thing forcing a re-upload every step.

Verified on a 7900 XTX (RADV): the original-dtype path matches torch.nn.functional.linear for M=1..8 across all three dtypes, and the test now covers M=0.

[ericcurtin]

Thanks for contributing @maci0 just poking you so you know this one needs work, didn't pass the build

[ericcurtin]

Review

Root cause and fix look correct: _vulkan_matvec's underlying shader tiles 4 rows, and padding+slicing in _wrap_linear is a reasonable dispatch-layer workaround. I checked the concerns that jumped out on first read and they turn out to be non-issues:

• Dropping the explicit weight.float() / relying on vulkan_ops.linear() to do the float conversion internally is fine — linear() (vulkan_ops.py:239,248,264) always calls .float() on x/weight/bias itself, so nothing is lost.
• xf[-1:].expand(4 - m, -1) produces a stride-0 view, but _to_bytes() calls .contiguous() before serializing, so that's safe too.

Blocking

CI is red: ruff I001 (unsorted import block) in tests/python/test_linear_dispatch.py:22-23 on all 3 lint jobs. Trivial — ruff check --fix or reorder the vulkan_ops/VulkanContext imports.

Needs coordination with #19

This PR and #19 both rewrite the same lines of vk_forward in vllm_vulkan/model_runner.py (the bias = getattr(...) / try-block / tuple-return section). I test-merged both locally and confirmed a real conflict:

git merge --no-commit --no-ff pr19   # on top of pr17
CONFLICT (content): Merge conflict in vllm_vulkan/model_runner.py

Neither PR is wrong on its own, but whichever merges second will need a manual rebase that keeps both the small-M padding and the TP collective logic together (i.e. pad-then-matmul-then-collective, not one replacing the other). Worth flagging to the author/reviewer merging these so the fix isn't silently dropped.

Mergeability: 8/10 — correct fix, tests included; blocked only by a one-line lint fix and coordinating the rebase with #19.

[ericcurtin]

Review (re-check against current main)

Conflicts

vllm_vulkan/model_runner.py conflicts with main, where vk_forward was rewritten (TP collectives, tuple return, weight-cache handling). Needs a rebase before anything else.

The fix is now obsolete

Main's vulkan_ops.linear (vulkan_ops.py:644-655) already dispatches on T:

• T < _MATVEC_THRESHOLD (= 4) goes to mul_mat_vec_{f16,f32}_f32_f32
• T >= 4 goes to the tiled mul_mm matmul

So the small-M tiled-matmul path this pads around is no longer reachable for decode. Padding M up to 4 would actively push decode into the tiled matmul: slower, and the exact path you measured as wrong.

The weight.float() to weight change also already landed on main (model_runner.py:376-396, with the cache-identity rationale documented).

Also

ruff I001 in tests/python/test_linear_dispatch.py:22-23, so lint is red on all three jobs.

Suggestion

Rebase and reduce this to just tests/python/test_linear_dispatch.py as a regression guard over the M = 0,1,2,3,4,8,17 dispatch boundary (it's genuinely useful, nothing on main covers M around the threshold), dropping the padding in model_runner.py. If the M<4 garbage still reproduces on current main, please reopen with that measurement instead; otherwise this can be closed.