A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 14:45 UTC
[ROCm][Bugfix] Fix ptpc scale load issue for fused shared expert path in deepseek mtp (#33018)
Signed-off-by: ganyi <ygan@amd.com>
Pleaplusone committed
Jan 26, 2026 at 23:19 UTC
be6931ee275b477dbd5e7cb1b7f99dcd37a6807a
1 file changed
+11
-8
vllm/model_executor/models/deepseek_mtp.py
+11
-8
@@ -316,7 +316,11 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
316
# Determine split axis based on op type
317
# gate/up: ColumnParallel → split along dim 0
318
# down: RowParallel → split along dim 1
319
- split_dim = 1 if "down_proj.weight" in name else 0
319
+ split_dim = (
320
+ 1
321
+ if ("down_proj.weight" in name and loaded_weight.ndim > 1)
322
+ else 0
323
+ )
324
total = loaded_weight.shape[split_dim]
325
assert total % num_chunks == 0, (
326
f"Shared expert weight dim {total} "
@@ -329,14 +333,13 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
333
weight_to_load = loaded_weight
334
335
if is_fusion_moe_shared_experts_layer:
332
- if split_dim == 0:
333
- weight_to_load = loaded_weight[
334
- j * chunk_size : (j + 1) * chunk_size, :
335
- ]
336
+ chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
337
+ if loaded_weight.ndim == 1:
338
+ weight_to_load = loaded_weight[chunk_slice]
339
+ elif split_dim == 0:
340
+ weight_to_load = loaded_weight[chunk_slice, :]
341
else:
337
- weight_to_load = loaded_weight[
338
- :, j * chunk_size : (j + 1) * chunk_size
339
- ]
342
+ weight_to_load = loaded_weight[:, chunk_slice]
343
# Synthesize an expert-style name so expert mapping
344
# can route it
345
chunk_name = name.replace(