修复awq-marlin的bug

09bdce60 · yangql · ef8dd155 · 09bdce60
Commit 09bdce60 authored Mar 03, 2026 by yangql
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

vllm/_custom_ops.py vllm/_custom_ops.py +3 -2

No files found.
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -1515,8 +1515,9 @@ def awq_marlin_moe_repack(
    output = torch.empty((num_experts, size_k // 16, size_n * (num_bits // 2)),
                         device=b_q_weight.device,
                         dtype=b_q_weight.dtype)
-    output[e] = op.awq_marlin_repack(b_q_weight[e], size_k,
-                                                   size_n, num_bits)
+    for e in range(num_experts):
+        output[e] = op.awq_marlin_repack(b_q_weight[e], size_k,
+                                                    size_n, num_bits)
    return output