[Bugfix] Fix MLA attention crash with AWQ/GPTQ quantized models

[haosdent]

Purpose

Fix AttributeError: 'ColumnParallelLinear' object has no attribute 'weight' when running MLA models with AWQ/GPTQ quantization (e.g., cyankiwi/GLM-4.7-Flash-AWQ-4bit).

Closes #34561.

Root cause: MLA attention code accesses self.kv_b_proj.weight.dtype in 3 places, but AWQ/GPTQ-quantized ColumnParallelLinear layers store weights as qweight (packed int32), not weight. The code only accounted for unquantized and FP8-quantized weights.

Fix: Guard .weight.dtype accesses with hasattr(self.kv_b_proj, "weight") checks:

• Line 406 (MLAAttention.__init__): Added hasattr guard in the and chain for the ROCm fp4 BMM check. Short-circuits to False for AWQ/GPTQ — correct, since packed int32 weights can't be used with fp4 BMM.
• Lines 2357-2371 (MLACommonImpl._compute_prefill_context): Introduced a local _kv_b_proj_w_dtype that uses weight.dtype when available, falling back to params_dtype (always present on LinearBase) for quantized layers. params_dtype is the model's compute dtype (e.g., bf16), which is the correct input dtype that AWQ/GPTQ layers expect.

Correctness verified across all quantization methods:

Quantization	hasattr(weight)	dtype used	fp4 BMM check	Prefill cast behavior
Unquantized bf16	True	bf16	Enabled if fp4bmm available	Cast to bf16
Unquantized fp16	True	fp16	Disabled (fp16≠bf16)	Cast to fp16
FP8	True	float8_e4m3fn	Disabled (fp8≠bf16)	fp8 prefill: cast to fp8; no fp8 prefill: skip cast
AWQ (W4A16)	False	params_dtype (bf16)	Disabled (short-circuit)	Cast to bf16
GPTQ (W4A16)	False	params_dtype (fp16)	Disabled (short-circuit)	Cast to fp16

Test Plan

1. Syntax and structural verification — AST parse, confirmed exactly 2 hasattr guards and 1 params_dtype fallback.
2. Mock-based logic verification — tested all 5 quantization scenarios (unquantized bf16, unquantized fp16, FP8 with/without FP8 prefill, AWQ) with mock objects simulating the ColumnParallelLinear layer behavior.

Test Result

Unquantized bf16: hasattr(weight)=True, weight.dtype=bf16
FP8: hasattr(weight)=True, weight.dtype=fp8
AWQ: hasattr(weight)=False, params_dtype=bf16
Line 406 logic: PASS for all 3 scenarios
Unquant, no fp8 prefill: cast=True, dtype=torch.bfloat16 -- PASS
FP8, no fp8 prefill: cast=False -- PASS
FP8, fp8 prefill: cast=True, dtype=torch.float8_e4m3fn -- PASS
AWQ, no fp8 prefill: cast=True, dtype=torch.bfloat16 -- PASS

All verification checks PASSED

[gemini-code-assist]

Code Review

This pull request effectively addresses a crash that occurs when running MLA models with AWQ/GPTQ quantization. The root cause, an AttributeError from accessing a non-existent .weight attribute on quantized layers, is correctly identified. The fix is clean and robust, using hasattr guards to prevent the error. The fallback to params_dtype for quantized layers in _compute_prefill_context is a logical and well-justified approach. The changes are minimal, targeted, and well-documented in the pull request description. Overall, this is an excellent bugfix.

[mgoin]

cc @LucasWilkinson @MatthewBonanni @pavanimajety to review

[cjackal]

I've checked that with this PR deepseek v3 awq checkpoint is loaded successfully and gets normal accuracy(gsm8k: 0.945) (naturally). The change looks accurate and simple, can we review this PR to unblock awq/gptq models from run on latest main w/ transformers v5?

[haosdent]

Thanks @cjackal 's test. @LucasWilkinson @MatthewBonanni @pavanimajety can you help to review? Thank you in advance.

[babyplutokurt]

I also verified this patch fixs the error serving GLM-4.7-Flash-GPTQ-4bits, with single and batch requests.

[MatthewBonanni]

LGTM, thanks for the contribution!