DeepSeek fix: awq x mergedreplicatedlinear

[mickaelseznec]

Purpose

Fixing #23530

Test Plan

Check manually that DeepSeek AWQ output is correct

Test Result

vllm (pretrained=/models/DeepSeek-R1-AWQ,tensor_parallel_size=8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9522	±	0.0059
		strict-match	5	exact_match	↑	0.9515	±	0.0059

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request aims to fix an issue with AWQ quantization in MergedReplicatedLinear layers by refactoring the weight loading to use a dedicated method, load_merged_column_weight. This is a good approach for handling specialized logic in custom parameter classes. However, the current implementation introduces a critical regression for unquantized models. It unconditionally calls param.load_merged_column_weight, but for unquantized layers, the parameter is a standard torch.nn.Parameter which lacks this method, and will cause an AttributeError during model loading. A check on the parameter type is needed to maintain backward compatibility for unquantized models.

[ashgold]

It would be great if you could assign a reviewer to get it applied quickly!

[cjackal]

I have also checked that this PR works for awq and fp8, thanks for the fix.

BTW we may not need the isinstance(param, PerTensorScaleParameter) check now?

[cjackal]

While I do think calling the kwarg tp_rank is misleading, it does the trick and I can't come up with a better naming.

[cjackal]

@mickaelseznec We can ping the corresponding reviewers by prepending DeepSeek to the title FYI.

[mickaelseznec]

Don't know who would be best to review, @mgoin @robertgshaw2-redhat because it's quantization related? Feel free to dispatch to others as well :)

[cjackal]

@tlrmchlsmth @yewentao256 This PR fixes weight loading logic of DeepSeek V2/V3 AWQ quantized models which is an oversight from the fused MLA qkv kernel update #21116.

[mgoin]

Why remove this case?

[cjackal]

PerTensorScaleParameter is a subclass of BasevLLMParameter, and PerTensorScaleParameter.load_merged_column_weight is doing the right weight loading, so I reckon it a code quality improvement.

TLDR: We have previously set shard_offset and shard_size to something not what the name represents just to make the following weight overloading part(param.data[shard_offset:shard_offset + shard_size] = loaded_weight) working, but PerTensorScaleParameter.load_merged_column_weight which is just BasevLLMParameter._assert_and_load is doing exactly the same.

[cjackal]

Can we get this PR reviewed and merged in 0.10.2? The content of the fix is not too complicated, if one needs e2e validation (beyond what is already given in the PR body) for the review I can help.

[qdivan]

LGTM

[Bulc]

Hello, I would like to know which version of vllm this PR will be merged into.

[xjpang]

Hello, I would like to know which version of vllm this PR will be merged into.

+1

[mickaelseznec]

Fix should already be there with #23024