[Core] Refactor QKVCrossParallelLinear implementation to support BNB 4-bit quantization

[Isotr0py]

• Revert [Bugfix] Revert QKVCrossParallelLinear usage in Mllama to keep BNB quantization work #14498 and add BNB support to QKVCrossParallelLinear

cc @NickLucche

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[NickLucche]

Thanks a lot for this!

I do agree the solution is a bit involved rn, but perhaps we can find a way to simplify it a bit..
Apart from the comments I left, I would also:

• add a test for unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit
• make sure VLLM_ATTENTION_BACKEND=XFORMERS python -m pytest -v tests/models/encoder_decoder/language/test_bart.py is also still working (not the case for me locally rn)

[Isotr0py]

add a test for unsloth/Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit

Mllama is quite large for testing, and we won't test it by running inference on CI. I think testing it on whisper or bart might be a better selection.

[NickLucche]

Sure but we already have tests for it, though they're guarded by the 48gb requirement so they won't run on L4, still it is useful to run them locally with a single command.

[jeejeelee]

LGTM, also cc @mgoin

[gshtras]

This PR has a regression, breaking support for amd/Llama-3.2-11B-Vision-Instruct-FP8-KV quantized models. @Isotr0py

[Isotr0py]

This PR has a regression, breaking support for amd/Llama-3.2-11B-Vision-Instruct-FP8-KV quantized models.

Oh, fp8 need to call process_weights_after_loading, let me think a method to handle this for QKVCrossParallelLinear.