@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f396bee56f

Fork 0

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 3, 22:45 UTC

[DSV4] Add PP support for deepseek-v4 (#41694)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: qizixi <22851944+zixi-qi@users.noreply.github.com>

Isotr0py committed May 10, 2026 at 23:47 UTC f396bee56fb516195a3669e5499129021ced4443

2 files changed +83 -22

docs/models/supported_models.md

+1 -1

@@ -385,7 +385,7 @@ th {

385 | `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |

386 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |

387 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |

390 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |

391 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |

vllm/model_executor/models/deepseek_v4.py

+82 -21

@@ -12,6 +12,7 @@ from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import VllmConfig, get_current_vllm_config

13 from vllm.distributed import (

14 get_ep_group,

15 + get_pp_group,

16 get_tensor_model_parallel_rank,

17 get_tensor_model_parallel_world_size,

18 )

@@ -49,6 +50,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

50 VocabParallelEmbedding,

51 )

52 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

53 +from vllm.model_executor.models.interfaces import SupportsPP

54 from vllm.model_executor.utils import set_weight_attrs

55 from vllm.platforms import current_platform

56 from vllm.sequence import IntermediateTensors

@@ -57,8 +59,10 @@ from vllm.utils.torch_utils import direct_register_custom_op

60 from .utils import (

61 AutoWeightsLoader,

62 + PPMissingLayer,

63 WeightsMapper,

64 extract_layer_index,

65 + is_pp_missing_parameter,

66 make_layers,

67 maybe_prefix,

68 )

@@ -1261,12 +1265,15 @@ class DeepseekV4Model(nn.Module):

1265 device=self.device,

1266 )

1267

1264 - self.embed_tokens = VocabParallelEmbedding(

1265 - config.vocab_size,

1266 - config.hidden_size,

1267 - quant_config=quant_config,

1268 - prefix=f"{prefix}.embed_tokens",

1269 - )

1268 + if get_pp_group().is_first_rank:

1269 + self.embed_tokens = VocabParallelEmbedding(

1270 + config.vocab_size,

1271 + config.hidden_size,

1272 + quant_config=quant_config,

1273 + prefix=f"{prefix}.embed_tokens",

1274 + )

1275 + else:

1276 + self.embed_tokens = PPMissingLayer()

1277

1278 self.start_layer, self.end_layer, self.layers = make_layers(

1279 config.num_hidden_layers,

@@ -1279,7 +1286,10 @@ class DeepseekV4Model(nn.Module):

1286 prefix=f"{prefix}.layers",

1287 )

1288

1282 - self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)

1289 + if get_pp_group().is_last_rank:

1290 + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)

1291 + else:

1292 + self.norm = PPMissingLayer()

1293

1294 self.hc_head_fn = nn.Parameter(

1295 torch.empty(

@@ -1304,16 +1314,42 @@ class DeepseekV4Model(nn.Module):

1314 # Pre-hc_head residual stream buffer for the MTP draft. Stable

1315 # address (outside the cudagraph pool) so the copy_ in forward()

1316 # refreshes it correctly across captured shapes.

1307 - self._mtp_hidden_buffer = torch.empty(

1308 - vllm_config.scheduler_config.max_num_batched_tokens,

1309 - self.hc_dim,

1310 - dtype=vllm_config.model_config.dtype,

1311 - device=self.device,

1312 - )

1317 + # refreshes it correctly across captured shapes. Only allocated on

1318 + # the last PP rank — that's where MTP target hidden states are

1319 + # produced.

1320 + if get_pp_group().is_last_rank:

1321 + self._mtp_hidden_buffer = torch.empty(

1322 + vllm_config.scheduler_config.max_num_batched_tokens,

1323 + self.hc_dim,

1324 + dtype=vllm_config.model_config.dtype,

1325 + device=self.device,

1326 + )

1327 + else:

1328 + self._mtp_hidden_buffer = None

1329

1330 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1331 return self.embed_tokens(input_ids)

1332

1333 + def make_empty_intermediate_tensors(

1334 + self,

1335 + batch_size: int,

1336 + dtype: torch.dtype,

1337 + device: torch.device,

1338 + ) -> IntermediateTensors:

1339 + # PP intermediate tensors carry the multi-stream hidden_states

1340 + # of shape (num_tokens, hc_mult, hidden_size) — V4 expands the

1341 + # token embedding to hc_mult streams before the first decoder

1342 + # layer and keeps that shape until hc_head() collapses it.

1343 + return IntermediateTensors(

1344 + {

1345 + "hidden_states": torch.zeros(

1346 + (batch_size, self.hc_mult, self.config.hidden_size),

1347 + dtype=dtype,

1348 + device=device,

1349 + ),

1350 + }

1351 + )

1352 +

1353 def forward(

1354 self,

1355 input_ids: torch.Tensor,

@@ -1321,8 +1357,16 @@ class DeepseekV4Model(nn.Module):

1357 intermediate_tensors: IntermediateTensors | None,

1358 inputs_embeds: torch.Tensor | None = None,

1359 ) -> torch.Tensor | IntermediateTensors:

1324 - hidden_states = self.embed_input_ids(input_ids)

1325 - hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)

1360 + if get_pp_group().is_first_rank:

1361 + if inputs_embeds is not None:

1362 + hidden_states = inputs_embeds

1363 + else:

1364 + hidden_states = self.embed_input_ids(input_ids)

1365 + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)

1366 + else:

1367 + assert intermediate_tensors is not None

1368 + hidden_states = intermediate_tensors["hidden_states"]

1369 +

1370 if self.use_mega_moe:

1371 input_ids = input_ids.to(torch.int64)

1372 for layer in islice(self.layers, self.start_layer, self.end_layer):

@@ -1332,6 +1376,9 @@ class DeepseekV4Model(nn.Module):

1376 input_ids,

1377 )

1378

1379 + if not get_pp_group().is_last_rank:

1380 + return IntermediateTensors({"hidden_states": hidden_states})

1381 +

1382 # Stash pre-hc_head residual for the MTP draft (captured copy_).

1383 num_tokens = hidden_states.shape[0]

1384 self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))

@@ -1380,6 +1427,8 @@ class DeepseekV4Model(nn.Module):

1427 continue

1428 name = name.replace(weight_name, param_name)

1429

1430 + if is_pp_missing_parameter(name, self):

1431 + break

1432 param = params_dict[name]

1433 weight_loader = param.weight_loader

1434 weight_loader(param, loaded_weight, shard_id)

@@ -1401,6 +1450,8 @@ class DeepseekV4Model(nn.Module):

1450 if weight_name not in name:

1451 continue

1452 name_mapped = name.replace(weight_name, param_name)

1453 + if is_pp_missing_parameter(name_mapped, self):

1454 + continue

1455 param = params_dict[name_mapped]

1456 # We should ask the weight loader to return success or not

1457 # here since otherwise we may skip experts with other

@@ -1422,12 +1473,16 @@ class DeepseekV4Model(nn.Module):

1473 loaded_params.add(name_mapped)

1474 continue

1475 elif "attn_sink" in name:

1476 + if is_pp_missing_parameter(name, self):

1477 + continue

1478 narrow_weight = loaded_weight[head_rank_start:head_rank_end]

1479 n = narrow_weight.shape[0]

1480 params_dict[name][:n].copy_(narrow_weight)

1481 loaded_params.add(name)

1482 continue

1483 else:

1484 + if is_pp_missing_parameter(name, self):

1485 + continue

1486 param = params_dict[name]

1487 weight_loader = getattr(

1488 param, "weight_loader", default_weight_loader

@@ -1525,7 +1580,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:

1580 )

1581

1582

1528 -class DeepseekV4ForCausalLM(nn.Module):

1583 +class DeepseekV4ForCausalLM(nn.Module, SupportsPP):

1584 model_cls = DeepseekV4Model

1585

1586 # Default mapper assumes the original FP4-expert checkpoint layout.

@@ -1544,12 +1599,18 @@ class DeepseekV4ForCausalLM(nn.Module):

1599 self.model = self.model_cls(

1600 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

1601 )

1547 - self.lm_head = ParallelLMHead(

1548 - config.vocab_size,

1549 - config.hidden_size,

1550 - prefix=maybe_prefix(prefix, "lm_head"),

1551 - )

1602 + if get_pp_group().is_last_rank:

1603 + self.lm_head = ParallelLMHead(

1604 + config.vocab_size,

1605 + config.hidden_size,

1606 + prefix=maybe_prefix(prefix, "lm_head"),

1607 + )

1608 + else:

1609 + self.lm_head = PPMissingLayer()

1610 self.logits_processor = LogitsProcessor(config.vocab_size)

1611 + self.make_empty_intermediate_tensors = (

1612 + self.model.make_empty_intermediate_tensors

1613 + )

1614

1615 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1616 return self.model.embed_input_ids(input_ids)