minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy (#31619)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-07-17 22:37:53 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 639261f7b2
commit 7a896215e7
7 changed files with 123 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/runner/eager_runner.py
+50 -48
View File
@@ -285,21 +285,7 @@ class EagerRunner(BaseRunner):
model_runner.attn_backend.init_forward_metadata(forward_batch)
cp_v2_active = is_cp_v2_active(forward_batch)
forward_positions = forward_batch.positions
if cp_v2_active:
prepare_cp_forward(forward_batch)
complete_hidden_states = kwargs.get("input_embeds")
if complete_hidden_states is None:
embed_layer = model_runner.model.get_input_embeddings()
complete_hidden_states = embed_layer(forward_batch.input_ids)
sharded_hidden_states, sharded_positions = cp_split_before_forward(
complete_hidden_states,
forward_batch.positions,
forward_batch,
)
kwargs["input_embeds"] = sharded_hidden_states
forward_positions = sharded_positions
else:
if not cp_v2_active:
forward_batch.attn_cp_metadata = None
category = (
@@ -335,51 +321,67 @@ class EagerRunner(BaseRunner):
):
ret = model_runner.model.forward(
forward_batch.input_ids,
forward_positions,
forward_batch.positions,
forward_batch,
**kwargs,
)
elif cp_v2_active:
# CP-V2: drive .model directly to gather across CP ranks before logits.
hidden_states = model_runner.model.model(
forward_batch.input_ids,
forward_positions,
forward_batch,
input_embeds=kwargs.get("input_embeds"),
pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
)
aux_hidden_states = None
capture_aux_hidden_states = getattr(
model_runner.model, "capture_aux_hidden_states", False
)
if capture_aux_hidden_states:
hidden_states, aux_hidden_states = hidden_states
if model_runner.model.pp_group.is_last_rank:
hidden_states = cp_gather_after_forward(
hidden_states,
forward_batch,
torch.cuda.current_stream(),
)
ret = model_runner.model.logits_processor(
forward_batch.input_ids,
hidden_states,
model_runner.model.lm_head,
forward_batch,
aux_hidden_states,
)
elif capture_aux_hidden_states:
ret = hidden_states, aux_hidden_states
else:
ret = hidden_states
ret = self._execute_extend_cp_v2(forward_batch, kwargs)
else:
ret = model_runner.model.forward(
forward_batch.input_ids,
forward_positions,
forward_batch.positions,
forward_batch,
**kwargs,
)
return ret
def _execute_extend_cp_v2(
self, forward_batch: ForwardBatch, kwargs: dict
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
"""CP-v2 extend: shard inputs at the model boundary, run the body on the
rank-local slice, then gather hidden states before the logits step.
"""
model = self.model_runner.model
prepare_cp_forward(forward_batch)
input_embeds = kwargs.get("input_embeds")
if input_embeds is None:
input_embeds = model.get_input_embeddings()(forward_batch.input_ids)
input_embeds, positions = cp_split_before_forward(
input_embeds, forward_batch.positions, forward_batch
)
hidden_states = model.model(
forward_batch.input_ids,
positions,
forward_batch,
input_embeds=input_embeds,
pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
)
capture_aux_hidden_states = getattr(model, "capture_aux_hidden_states", False)
aux_hidden_states = None
if capture_aux_hidden_states:
hidden_states, aux_hidden_states = hidden_states
if not model.pp_group.is_last_rank:
return (
(hidden_states, aux_hidden_states)
if capture_aux_hidden_states
else hidden_states
)
hidden_states = cp_gather_after_forward(
hidden_states, forward_batch, torch.cuda.current_stream()
)
return model.logits_processor(
forward_batch.input_ids,
hidden_states,
model.lm_head,
forward_batch,
aux_hidden_states,
)
def _execute_idle(
self, forward_batch: ForwardBatch, pp_proxy_tensors=None
) -> Union[LogitsProcessorOutput, PPProxyTensors]: