Fix AWQ checkpoint loading for Qwen3.5 fused GDN QKV projection

[adityaanikam]

Motivation

Fixes #4899.

Loading an AWQ checkpoint for a Qwen3.5 hybrid-GDN model with --backend pytorch --tp 2 kills the engine subprocess during weight loading:

File ".../lmdeploy/pytorch/models/qwen3_5.py", line 503, in qkv_weight_loader
    bs = mod.block_size
AttributeError: 'AwqLinear' object has no attribute 'block_size'

Qwen3_5GatedDeltaNet._patch_qkv_weight_loader patches weight_loader to do the non-uniform [key_dim, key_dim, value_dim] TP split. It was written for the default/fp8 layouts, where the output features live on dim0 and a scale tensor is dim0-shrunk by block_size. AwqLinear has neither block_size nor that layout, so the scale branch raises.

Fixing only the attribute would leave a second, quieter problem. AWQ keeps output features on the last dim: qweight is (in_features, out_features // elem_per_int), scales is (in_features // group_size, out_features), qzeros matches qweight. With key_dim=512, value_dim=1024, in_features=2048, qweight is (2048, 256) — the shape[0] < sum(sections) heuristic is False, so the old code splits it along dim0, i.e. along in_features. That is the wrong axis and corrupts weights silently rather than raising.

Modification

Add an AWQ branch to qkv_weight_loader that splits along the last dim and scales the sections by the packing factor for the int32-packed tensors, following the existing MergedAwqLinear.weight_loader convention:

• scales / bias -> split by the raw sections
• qweight / qzeros -> split by section // elem_per_int
• chunk along dim=-1

AWQ is detected via hasattr(mod, 'elem_per_int'). That is mutually exclusive with the fp8 path: AwqLinear defines elem_per_int and group_size but no block_size, and the blocked-fp8 linear defines block_size but neither of the others. The existing non-AWQ code path is untouched.

BC-breaking (Optional)

No. Only adds a branch for AWQ modules, which previously raised AttributeError on this path.

Use cases (Optional)

Serving an AWQ (w4a16) Qwen3.5 hybrid-GDN checkpoint on the PyTorch engine with tp > 1. tp=1 was unaffected (mod.is_tp == False short-circuits to default_weight_loader).

Checklist

1. ruff check passes on the modified file. (ruff format was not run, as it would reformat unrelated pre-existing code; pre-commit runs ruff-check.)
2. Verified without GPU by driving the patched qkv_weight_loader with stub modules matching real AwqLinear / blocked-fp8 attributes and tensor shapes: all four AWQ parameter types produce correct per-rank shard shapes and reassemble exactly per section at tp=2; the fp8 path produces identical output to before; and reverting the change reproduces the reported AttributeError. I do not have the reporter's hardware or checkpoint, so a confirmation run on the failing artifact would be welcome — the reporter offered one in the issue.
3. No downstream version dependency.
4. No documentation change needed.

[Copilot]

🟡 Changes recommended

The new AWQ sharding path uses Tensor.chunk() instead of AWQ’s alignment-aware partitioning, which can produce shard shapes that don’t match AwqLinear TP sizing for some valid configurations.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR fixes AWQ checkpoint loading for Qwen3.5 hybrid-GDN models when using PyTorch backend with tensor-parallelism (tp > 1) by updating the patched fused QKV weight_loader to correctly shard AWQ-packed tensors along the output (last) dimension.

Changes:

• Add an AWQ-specific branch in qkv_weight_loader to split qweight/qzeros using elem_per_int-scaled sections and to split along dim=-1.
• Avoid the existing fp8/block-size heuristic on AWQ modules (which previously raised AttributeError: ... block_size and could also silently shard on the wrong axis).

File summaries

File	Description
lmdeploy/pytorch/models/qwen3_5.py	Adds AWQ-aware sharding logic for the fused GDN QKV projection loader under TP.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[qescccczmr]

Could you add a TP sharding regression test that checks both shard contents and destination parameter shapes for supported Qwen3.5 configurations? This branch uses chunk, whereas the generic AWQ loaders use chunk_aligned; the test would help document why equal per-section partitioning is appropriate here.

[adityaanikam]

Added test_awq_qkv_shard_matches_chunk_aligned in tests/pytorch/test_qwen3_5.py. It checks both shard contents and destination shapes at tp=2 and tp=4 for the #4899 configuration, using chunk_aligned itself as the reference for these dimensions each packed section is an exact multiple of world_size * elem_per_int, so chunk_aligned has no remainder to redistribute and matches plain chunk exactly. Added that reasoning as a module docstring too.