[ROCm][Model][Bugfix] Fix GLM-5.2 shared-expert fusion and MTP on the ROCm DSA path (#58904)
Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
Jin Tao committed
Sep 29, 2026 at 10:57 UTC
77d9cda688a21087539a495562134a6f2f205efe
2 files changed
+106
-36
vllm/models/deepseek_v32/amd/model.py
+97
-36
@@ -12,6 +12,9 @@ from vllm.distributed import get_pp_group
12
from vllm.model_executor.layers.fused_moe import (
13
fused_moe_make_expert_params_mapping,
14
)
15
+from vllm.model_executor.layers.fused_moe.utils import (
16
+ is_model_fused_shared_expert_compatible,
17
+)
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.vocab_parallel_embedding import (
20
VocabParallelEmbedding,
@@ -158,6 +161,12 @@ class DeepseekV32Model(torch.nn.Module):
161
prefix=f"{prefix}.layers",
162
)
163
164
+ self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
165
+ self.layers,
166
+ DeepseekV2MoE,
167
+ "mlp",
168
+ )
169
+
170
if get_pp_group().is_last_rank:
171
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
172
else:
@@ -226,7 +235,12 @@ class DeepseekV32Model(torch.nn.Module):
235
ckpt_gate_proj_name="gate_proj",
236
ckpt_down_proj_name="down_proj",
237
ckpt_up_proj_name="up_proj",
229
- num_experts=self.config.n_routed_experts,
238
+ num_experts=self.config.n_routed_experts
239
+ + (
240
+ self.config.n_shared_experts
241
+ if self.is_fused_shared_expert_enabled
242
+ else 0
243
+ ),
244
num_redundant_experts=self.num_redundant_experts,
245
)
246
@@ -239,6 +253,9 @@ class DeepseekV32Model(torch.nn.Module):
253
continue
254
if get_spec_layer_idx_from_weight_name(self.config, name) is not None:
255
continue
256
+ is_fusion_moe_shared_experts_layer = (
257
+ self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)
258
+ )
259
if _try_load_fp8_indexer_wk(
260
name,
261
loaded_weight,
@@ -254,6 +271,8 @@ class DeepseekV32Model(torch.nn.Module):
271
continue
272
if ("mlp.experts." in name) and name not in params_dict:
273
continue
274
+ if is_fusion_moe_shared_experts_layer:
275
+ continue
276
name_mapped = name.replace(weight_name, param_name)
277
if (
278
param_name == "fused_qkv_a_proj"
@@ -269,43 +288,85 @@ class DeepseekV32Model(torch.nn.Module):
288
break
289
else:
290
is_expert_weight = False
272
- for mapping in expert_params_mapping:
273
- param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
274
- if weight_name not in name:
275
- continue
276
- is_expert_weight = True
277
- name_mapped = name.replace(weight_name, param_name)
278
- if is_pp_missing_parameter(name_mapped, self):
279
- continue
280
- param = params_dict[name_mapped]
281
- weight_loader = typing.cast(
282
- Callable[..., bool], param.weight_loader
291
+
292
+ # When AITER fusion_shared_experts is enabled the checkpoint
293
+ # provides a single widened shared_experts tensor per layer with
294
+ # no expert index. Split it evenly and route the slices into the
295
+ # expert slots appended after the routed experts.
296
+ num_chunks = 1
297
+ if is_fusion_moe_shared_experts_layer:
298
+ num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
299
+ # gate/up are ColumnParallel (dim 0), down is RowParallel (dim 1).
300
+ split_dim = (
301
+ 1
302
+ if ("down_proj.weight" in name and loaded_weight.ndim > 1)
303
+ else 0
304
)
284
- success = weight_loader(
285
- param,
286
- loaded_weight,
287
- name_mapped,
288
- shard_id=shard_id,
289
- expert_id=expert_id,
290
- return_success=True,
305
+ total = loaded_weight.shape[split_dim]
306
+ assert total % num_chunks == 0, (
307
+ f"Shared expert weight dim {total} "
308
+ f"not divisible by num_chunks {num_chunks}"
309
)
292
- if success:
293
- name = name_mapped
294
- break
295
- else:
296
- if is_expert_weight:
297
- continue
298
- if name.endswith(".bias") and name not in params_dict:
299
- continue
300
- name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
301
- if name is None:
302
- continue
303
- if is_pp_missing_parameter(name, self):
304
- continue
305
- param = params_dict[name]
306
- loader = getattr(param, "weight_loader", default_weight_loader)
307
- loader(param, loaded_weight)
308
- loaded_params.add(name)
310
+ chunk_size = total // num_chunks
311
+
312
+ for j in range(num_chunks):
313
+ chunk_name = name
314
+ weight_to_load = loaded_weight
315
+
316
+ if is_fusion_moe_shared_experts_layer:
317
+ chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
318
+ if loaded_weight.ndim == 1:
319
+ weight_to_load = loaded_weight[chunk_slice]
320
+ elif split_dim == 0:
321
+ weight_to_load = loaded_weight[chunk_slice, :]
322
+ else:
323
+ weight_to_load = loaded_weight[:, chunk_slice]
324
+ chunk_name = name.replace(
325
+ "mlp.shared_experts",
326
+ f"mlp.experts.{self.config.n_routed_experts + j}",
327
+ )
328
+
329
+ for mapping in expert_params_mapping:
330
+ param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
331
+ if weight_name not in chunk_name:
332
+ continue
333
+ is_expert_weight = True
334
+ name_mapped = chunk_name.replace(weight_name, param_name)
335
+ if is_pp_missing_parameter(name_mapped, self):
336
+ continue
337
+ param = params_dict[name_mapped]
338
+ weight_loader = typing.cast(
339
+ Callable[..., bool], param.weight_loader
340
+ )
341
+ success = weight_loader(
342
+ param,
343
+ weight_to_load,
344
+ name_mapped,
345
+ shard_id=shard_id,
346
+ expert_id=expert_id,
347
+ return_success=True,
348
+ )
349
+ if success:
350
+ if not is_fusion_moe_shared_experts_layer:
351
+ name = name_mapped
352
+ else:
353
+ loaded_params.add(name_mapped)
354
+ break
355
+ else:
356
+ if is_expert_weight:
357
+ continue
358
+ if name.endswith(".bias") and name not in params_dict:
359
+ continue
360
+ name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
361
+ if name is None:
362
+ continue
363
+ if is_pp_missing_parameter(name, self):
364
+ continue
365
+ param = params_dict[name]
366
+ loader = getattr(param, "weight_loader", default_weight_loader)
367
+ loader(param, loaded_weight)
368
+ if name is not None and not is_fusion_moe_shared_experts_layer:
369
+ loaded_params.add(name)
370
return loaded_params
371
372
vllm/models/deepseek_v32/amd/mtp.py
+9
@@ -130,6 +130,15 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
130
if self_attn is not None and hasattr(self_attn, "skip_topk"):
131
self_attn.skip_topk = skip
132
133
+ def compact_topk_indices(self, slot_ids: torch.Tensor):
134
+ """Gather the top-k index rows at ``slot_ids`` to the front of the buffer."""
135
+ num_slots = slot_ids.numel()
136
+ for layer in self.layers.values():
137
+ self_attn = getattr(layer.mtp_block, "self_attn", None)
138
+ if self_attn is not None and hasattr(self_attn, "topk_indices_buffer"):
139
+ topk_indices_buffer = self_attn.topk_indices_buffer
140
+ topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
141
+
142
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
143
return self.embed_tokens(input_ids)
144