[K3 Perf] Combine multiple all gather together for SP, 1.5~3x kernel level performance improvement (#51070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 5, 2026 at 17:42 UTC
877975897d68f5e837e25d2f8ba80f5355adbe5b
1 file changed
+21
-15
vllm/models/kimi_k3/nvidia/model.py
+21
-15
@@ -1118,13 +1118,6 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1118
residual = intermediate_tensors["residual"]
1119
assert hidden_states is not None
1120
1121
- aux_hidden_states: list[torch.Tensor] = []
1122
- if self.start_layer in self.aux_hidden_state_layers:
1123
- if self.use_attn_res or residual is None:
1124
- aux_hidden_states.append(hidden_states)
1125
- else:
1126
- aux_hidden_states.append(hidden_states + residual)
1127
-
1121
full_num_tokens = positions.shape[0]
1122
if self.use_sequence_parallel:
1123
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
@@ -1135,6 +1128,14 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1128
hidden_states = sp_shard(hidden_states)
1129
assert residual is None, "Currently, SP is not supported with PP"
1130
1131
+ # sharded aux hidden states when sp is enabled
1132
+ aux_hidden_states: list[torch.Tensor] = []
1133
+ if self.start_layer in self.aux_hidden_state_layers:
1134
+ if self.use_attn_res or residual is None:
1135
+ aux_hidden_states.append(hidden_states)
1136
+ else:
1137
+ aux_hidden_states.append(hidden_states + residual)
1138
+
1139
prefix_sum = None
1140
if self.use_attn_res:
1141
block_residual = hidden_states.new_empty(
@@ -1166,11 +1167,6 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1167
assert residual is not None
1168
aux_hidden_state = hidden_states + residual
1169
1169
- if self.use_sequence_parallel:
1170
- # Gather SP-sharded aux hidden states.
1171
- # TODO: Optimize this.
1172
- aux_hidden_state = sp_all_gather(aux_hidden_state)
1173
- aux_hidden_state = aux_hidden_state[:full_num_tokens]
1170
aux_hidden_states.append(aux_hidden_state)
1171
1172
assert hidden_states is not None
@@ -1203,9 +1199,19 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1199
hidden_states = hidden_states + residual
1200
1201
if self.use_sequence_parallel:
1206
- # Gather SP-sharded hidden states.
1207
- hidden_states = sp_all_gather(hidden_states)
1208
- hidden_states = hidden_states[:full_num_tokens]
1202
+ if aux_hidden_states:
1203
+ hidden_size = hidden_states.shape[-1]
1204
+ packed_hidden_states = torch.cat(
1205
+ [hidden_states, *aux_hidden_states], dim=-1
1206
+ )
1207
+ packed_hidden_states = sp_all_gather(packed_hidden_states)
1208
+ packed_hidden_states = packed_hidden_states[:full_num_tokens]
1209
+ hidden_states, *aux_hidden_states = packed_hidden_states.split(
1210
+ hidden_size, dim=-1
1211
+ )
1212
+ else:
1213
+ hidden_states = sp_all_gather(hidden_states)
1214
+ hidden_states = hidden_states[:full_num_tokens]
1215
1216
# NOTE: the final norm is applied in compute_logits instead of here, so
1217
# the MTP draft model receives the pre-norm hidden states.