[Bugfix] Handle process_weights_after_loading for QKVCrossParallelLinear

[Isotr0py]

FIX #14545 (comment)

[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.

🚀

[Isotr0py]

@gshtras Can you check if this PR can fix the Mllama FP8 regression issue? I can generate reasonable outputs with neuralmagic/Llama-3.2-11B-Vision-Instruct-FP8-dynamic, but not amd/Llama-3.2-11B-Vision-Instruct-FP8-KV.

I'm not sure if it's due to hardware difference or something other, because I can't generate reasonable outputs from this checkpoint with an earlier commit (ec79b67) before QKVCrossParallelLinear introduction as well.

[NickLucche]

Thanks for the fix!

[gshtras]

@gshtras Can you check if this PR can fix the Mllama FP8 regression issue? I can generate reasonable outputs with neuralmagic/Llama-3.2-11B-Vision-Instruct-FP8-dynamic, but not amd/Llama-3.2-11B-Vision-Instruct-FP8-KV.

I'm not sure if it's due to hardware difference or something other, because I can't generate reasonable outputs from this checkpoint with an earlier commit (ec79b67) before QKVCrossParallelLinear introduction as well.

This fixes the weight loading, but there are also runtime issues in attention now, since 77a318b
Not yet sure if it is related, or a new issue

[Isotr0py]

there are also runtime issues in attention now, since 77a318b

Hmmm, that's weird. I can run this model without any errors, but the problem is that the output hidden states will contain NaN, causing model outputs all "!".

[gshtras]

there are also runtime issues in attention now, since 77a318b

Hmmm, that's weird. I can run this model without any errors, but the problem is that the output hidden states will contain NaN, causing model outputs all "!".

Possibly hardware related, to clarify, I'm checking with V0 on a MI300X machine.
Applying this fix to the ROCm fork solves the issue, so there appears to be an unrelated regression, I think this PR is GTG, thanks

[gshtras]

there are also runtime issues in attention now, since 77a318b

Hmmm, that's weird. I can run this model without any errors, but the problem is that the output hidden states will contain NaN, causing model outputs all "!".

So that is due to an issue with the graph mode, which shouldn't be used on ROCm and is addressed in #15413

[SageMoore]

I generally think this looks fine, but I'd like someone with more experience in this part of the codebase to take a look. @robertgshaw2-redhat @tlrmchlsmth @mgoin