[Bugfix][ROCm] Fixing the skinny gemm dispatch logic from #32831

[gshtras]

Fixing 2 issues from the previous PR

1. When the wvsplitk kernel is not applicable, the fallback should be hipblaslt through torch._scaled_mm, and not cloning and modifying the input tensor. This fixes the +100% performance regression on amd/LLama-3.1-*-FP8-KV models when running on smaller batch sizes
2. Weights can be non-contiguous, and often are in the FP8 case where we explicitly pad them to a multiple of 256. So the condition needs to only be applied to the activation tensor
3. cc @rasmith

[gemini-code-assist]

Code Review

This pull request correctly refactors the skinny GEMM dispatch logic by removing the unconditional enforcement of contiguous inputs and instead adding checks at the call sites. For the FP8 kernel wvSplitKQ, checking only the activation tensor's contiguity is correct as the kernel can handle non-contiguous weights. However, for the unquantized kernels (wvSplitK, wvSplitKrc, LLMM1), they still seem to require both inputs to be contiguous. The current changes only check the activation tensor, which could lead to incorrect results if non-contiguous weights are passed. I've added comments to vllm/model_executor/layers/utils.py to suggest adding checks for weight contiguity for these unquantized kernels to prevent this potential issue.

[JartX]

@gshtras Excuse me for asking, but could this PR break inference on gfx1x by removing part of the contiguous logic? There have been several instances where removing contiguous logic caused the model specifically in ViT to produce nothing but hallucinations. Thank you for your time.

[gshtras]

@gshtras Excuse me for asking, but could this PR break inference on gfx1x by removing part of the contiguous logic? There have been several instances where removing contiguous logic caused the model specifically in ViT to produce nothing but hallucinations. Thank you for your time.

No, the condition to not dispatch these kernels on Radeon is unchanged

[JartX]

@gshtras many thanks for the answer

[rasmith]

@gshtras The tests that were breaking were from

pytest -sv models/language/generation/test_hybrid.py

and I tested them and found no failures.

Also, I tested reverting https://github.com/vllm-project/vllm/pull/32099 with this PR applied and also found no failures in the test_hybrid.py test.

Could you also revert that PR as well in this one?

[tjtanaa]

I also help to test with the command used in PR #32831

vllm bench serve --model state-spaces/mamba-130m-hf
vllm serve state-spaces/mamba-130m-hf

The performance seems good. Fallback is better than just trying to cast to contiguous for the model as well.

[tjtanaa]

LGTM. Thanks @gshtras for the quick fix.

[tjtanaa]

@gshtras can you try rebasing again? I have retried those tests and they still failed.