Skip to content

Commit b7f7654

Browse files
[BugFix][310p] Qwen3.5 shared-experts linear failuer when out_dim=1 (vllm-project#11714)
### What this PR does / why we need it? This PR fixes a failure in Qwen3.5 shared-experts linear layer on 310P devices when `out_dim=1`. On 310P, `torch_npu` rejects `FRACTAL_NZ` matmul when the weight-side matrix has $N=1$ or $K=1$. To resolve this, we introduce `_should_keep_nd_for_310p_weight` to detect if the weight has a dimension of 1 on 310P, and if so, keep the weight in ND format instead of converting it to `FRACTAL_NZ`. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Tested on 310P Ascend NPU. - vLLM version: v0.23.0 - vLLM main: vllm-project/vllm@1f486d9 Signed-off-by: Tflowers-0129 <2906339855@qq.com> Co-authored-by: Yizhou <136800916+yiz-liu@users.noreply.github.com>
1 parent c062a7e commit b7f7654

1 file changed

Lines changed: 13 additions & 2 deletions

File tree

vllm_ascend/ops/linear.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
AscendDeviceType,
4646
enable_sp,
4747
get_ascend_device_type,
48+
is_310p,
4849
maybe_trans_nz,
4950
)
5051

@@ -75,16 +76,26 @@ def unquantized_gemm_fake(
7576
)
7677

7778

79+
def _should_keep_nd_for_310p_weight(weight: torch.Tensor) -> bool:
80+
return is_310p() and weight.ndim >= 2 and (weight.shape[-1] == 1 or weight.shape[-2] == 1)
81+
82+
7883
class AscendUnquantizedLinearMethod(UnquantizedLinearMethod):
7984
"""Linear method without quantization"""
8085

8186
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
8287
super().process_weights_after_loading(layer)
88+
keep_nd_weight = _should_keep_nd_for_310p_weight(layer.weight.data)
8389
# must use fp32 to avoid accuracy degradation in dsv4.
8490
if getattr(layer, "precast_fp32_weight", False):
85-
layer.weight_fp32 = maybe_trans_nz(layer.weight.data.to(torch.float32))
91+
weight_fp32 = layer.weight.data.to(torch.float32)
92+
layer.weight_fp32 = weight_fp32 if keep_nd_weight else maybe_trans_nz(weight_fp32)
8693
if "conv1d" not in layer.prefix:
87-
layer.weight.data = maybe_trans_nz(layer.weight.data)
94+
# 310P torch_npu rejects FRACTAL_NZ matmul when the weight-side
95+
# matrix has n=1 or k=1. Keep scalar gates such as Qwen MoE's
96+
# shared_expert_gate in ND format, leaving non-310P policy intact.
97+
if not keep_nd_weight:
98+
layer.weight.data = maybe_trans_nz(layer.weight.data)
8899

89100
def apply(
90101
self,

0 commit comments

Comments
 (0)