5c91039c41 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -42,6 +42,7 @@ from vllm.distributed import (

42 get_tensor_model_parallel_rank,

43 get_tensor_model_parallel_world_size,

44 tensor_model_parallel_all_gather,

45 + tensor_model_parallel_reduce_scatter,

46 )

47 from vllm.logger import init_logger

48 from vllm.model_executor.layers.activation import SiluAndMul

@@ -128,6 +129,7 @@ class DeepseekAttention(nn.Module):

129 max_position_embeddings: int = 8192,

130 cache_config: CacheConfig | None = None,

131 quant_config: QuantizationConfig | None = None,

132 + reduce_results: bool = True,

133 prefix: str = "",

134 **kwargs,

135 ) -> None:

@@ -166,6 +168,7 @@ class DeepseekAttention(nn.Module):

168 self.total_num_heads * self.head_dim,

169 hidden_size,

170 bias=False,

171 + reduce_results=reduce_results,

172 quant_config=quant_config,

173 )

174

@@ -372,15 +375,17 @@ class DeepseekV2MoE(nn.Module):

375 self.gate.e_score_correction_bias.data.to(self.gate.out_dtype)

376 )

377

375 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

378 + def forward(

379 + self,

380 + hidden_states: torch.Tensor,

381 + already_sequence_parallel: bool = False,

382 + ) -> torch.Tensor:

383 num_tokens, hidden_dim = hidden_states.shape

384 hidden_states = hidden_states.view(-1, hidden_dim)

385

386 # Chunk the hidden states so they aren't replicated across TP ranks.

387 # This avoids duplicate computation in self.experts.

381 - # TODO: We can replace the all_reduce at the end of attn with a

382 - # reduce_scatter instead of chunking here.

383 - if self.is_sequence_parallel:

388 + if self.is_sequence_parallel and not already_sequence_parallel:

389 hidden_states = sequence_parallel_chunk(hidden_states)

390

391 if self.experts.is_internal_router:

@@ -393,7 +398,7 @@ class DeepseekV2MoE(nn.Module):

398 hidden_states=hidden_states, router_logits=router_logits

399 )

400

396 - if self.is_sequence_parallel:

401 + if self.is_sequence_parallel and not already_sequence_parallel:

402 final_hidden_states = tensor_model_parallel_all_gather(

403 final_hidden_states, 0

404 )

@@ -436,6 +441,7 @@ class DeepseekV2Attention(nn.Module):

441 cache_config: CacheConfig | None = None,

442 quant_config: QuantizationConfig | None = None,

443 topk_indices_buffer: torch.Tensor | None = None,

444 + reduce_results: bool = True,

445 prefix: str = "",

446 ) -> None:

447 super().__init__()

@@ -502,6 +508,7 @@ class DeepseekV2Attention(nn.Module):

508 self.num_heads * self.v_head_dim,

509 self.hidden_size,

510 bias=False,

511 + reduce_results=reduce_results,

512 quant_config=quant_config,

513 prefix=f"{prefix}.o_proj",

514 )

@@ -950,6 +957,7 @@ class DeepseekV2MLAAttention(nn.Module):

957 prefix: str = "",

958 topk_indices_buffer: torch.Tensor | None = None,

959 input_size: int | None = None,

960 + reduce_results: bool = True,

961 ) -> None:

962 super().__init__()

963 self.hidden_size = hidden_size

@@ -1018,6 +1026,7 @@ class DeepseekV2MLAAttention(nn.Module):

1026 self.num_heads * self.v_head_dim,

1027 self.hidden_size,

1028 bias=False,

1029 + reduce_results=reduce_results,

1030 quant_config=quant_config,

1031 prefix=f"{prefix}.o_proj",

1032 )

@@ -1184,6 +1193,18 @@ class DeepseekV2DecoderLayer(nn.Module):

1193 attn_cls = DeepseekV2MLAAttention

1194 else:

1195 attn_cls = DeepseekV2Attention

1196 + is_moe_layer = (

1197 + config.n_routed_experts is not None

1198 + and layer_idx >= config.first_k_dense_replace

1199 + and layer_idx % moe_layer_freq == 0

1200 + )

1201 + # TODO(wentao): enable SP MoE with PP after the PP boundary logic can safely

1202 + # send/receive sequence-parallel hidden_states across stages.

1203 + self.use_sequence_parallel_moe = (

1204 + parallel_config.use_sequence_parallel_moe

1205 + and parallel_config.pipeline_parallel_size == 1

1206 + and is_moe_layer

1207 + )

1208 self.self_attn = attn_cls(

1209 vllm_config=vllm_config,

1210 config=config,

@@ -1199,13 +1220,10 @@ class DeepseekV2DecoderLayer(nn.Module):

1220 quant_config=quant_config,

1221 prefix=f"{prefix}.self_attn",

1222 topk_indices_buffer=topk_indices_buffer,

1223 + reduce_results=not self.use_sequence_parallel_moe,

1224 )

1225

1204 - if (

1205 - config.n_routed_experts is not None

1206 - and layer_idx >= config.first_k_dense_replace

1207 - and layer_idx % moe_layer_freq == 0

1208 - ):

1226 + if is_moe_layer:

1227 self.mlp = DeepseekV2MoE(

1228 config=config,

1229 parallel_config=parallel_config,

@@ -1233,6 +1251,13 @@ class DeepseekV2DecoderLayer(nn.Module):

1251 residual: torch.Tensor | None,

1252 llama_4_scaling: torch.Tensor | None = None,

1253 ) -> torch.Tensor:

1254 + full_num_tokens = positions.shape[0]

1255 + input_is_sequence_parallel = (

1256 + self.use_sequence_parallel_moe

1257 + and residual is not None

1258 + and hidden_states.shape[0] != full_num_tokens

1259 + )

1260 +

1261 # Self Attention

1262 if residual is None:

1263 residual = hidden_states

@@ -1240,6 +1265,10 @@ class DeepseekV2DecoderLayer(nn.Module):

1265 else:

1266 hidden_states, residual = self.input_layernorm(hidden_states, residual)

1267

1268 + if input_is_sequence_parallel:

1269 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

1270 + hidden_states = hidden_states[:full_num_tokens]

1271 +

1272 attn_kwargs = {

1273 "positions": positions,

1274 "hidden_states": hidden_states,

@@ -1261,9 +1290,29 @@ class DeepseekV2DecoderLayer(nn.Module):

1290 # first layer.

1291 residual *= 1.0 / self.routed_scaling_factor

1292

1293 + if self.use_sequence_parallel_moe:

1294 + sp_remainder = (

1295 + hidden_states.shape[0] % get_tensor_model_parallel_world_size()

1296 + )

1297 + # pad if not divisible by world size

1298 + if sp_remainder:

1299 + sp_pad = get_tensor_model_parallel_world_size() - sp_remainder

1300 + hidden_states = torch.nn.functional.pad(

1301 + hidden_states, (0, 0, 0, sp_pad)

1302 + )

1303 + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)

1304 + if not input_is_sequence_parallel:

1305 + residual = sequence_parallel_chunk(residual)

1306 +

1307 # Fully Connected

1308 hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

1266 - hidden_states = self.mlp(hidden_states)

1309 + if self.use_sequence_parallel_moe:

1310 + hidden_states = self.mlp(

1311 + hidden_states,

1312 + already_sequence_parallel=True,

1313 + )

1314 + else:

1315 + hidden_states = self.mlp(hidden_states)

1316

1317 if isinstance(self.mlp, DeepseekV2MLP) and hidden_states.dtype == torch.float16:

1318 # Fix FP16 overflow

@@ -1385,8 +1434,25 @@ class DeepseekV2Model(nn.Module):

1434 islice(self.layers, self.start_layer, self.end_layer),

1435 start=self.start_layer,

1436 ):

1437 + # all gather if we need to use the whole states

1438 + if (

1439 + hidden_states.shape[0] != positions.shape[0]

1440 + and not layer.use_sequence_parallel_moe

1441 + ):

1442 + combined_states = torch.cat([hidden_states, residual], dim=-1)

1443 + combined_states = tensor_model_parallel_all_gather(combined_states, 0)

1444 + combined_states = combined_states[: positions.shape[0]]

1445 + hidden_states, residual = combined_states.split(

1446 + [self.hidden_size, self.hidden_size], dim=-1

1447 + )

1448 if idx in self.aux_hidden_state_layers:

1389 - aux_hidden_states.append(hidden_states + residual)

1449 + aux_hidden_state = hidden_states + residual

1450 + if aux_hidden_state.shape[0] != positions.shape[0]:

1451 + aux_hidden_state = tensor_model_parallel_all_gather(

1452 + aux_hidden_state, 0

1453 + )

1454 + aux_hidden_state = aux_hidden_state[: positions.shape[0]]

1455 + aux_hidden_states.append(aux_hidden_state)

1456 hidden_states, residual = layer(

1457 positions, hidden_states, residual, llama_4_scaling

1458 )

@@ -1396,6 +1462,14 @@ class DeepseekV2Model(nn.Module):

1462 {"hidden_states": hidden_states, "residual": residual}

1463 )

1464

1465 + if hidden_states.shape[0] != positions.shape[0]:

1466 + combined_states = torch.cat([hidden_states, residual], dim=-1)

1467 + combined_states = tensor_model_parallel_all_gather(combined_states, 0)

1468 + combined_states = combined_states[: positions.shape[0]]

1469 + hidden_states, residual = combined_states.split(

1470 + [self.hidden_size, self.hidden_size], dim=-1

1471 + )

1472 +

1473 hidden_states, _ = self.norm(hidden_states, residual)

1474 if len(aux_hidden_states) > 0:

1475 return hidden_states, aux_hidden_states

[GLM5.2 Perf] Replace MOE all-reduce with reduce-scatter, 3.1%~3.2 E2E Throughput improvement (#46635)