chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Revert "[SGL] sync patch: Remove sync points, prefill cudagraph for DP, disable cache reset in mem check (#19190)" (#19581)

Browse Source 
Co-authored-by: Alison Shao <alisonshao@mac.lan>
This commit is contained in:
Alison Shao
2026-02-28 19:46:47 -08:00
committed by GitHub
parent e64095c3c7
commit a45613f2a6
7 changed files with 42 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/layers/logits_processor.py
View File

@@ -519,11 +519,11 @@ class LogitsProcessor(nn.Module):
if hidden_states_before_norm is not None:
pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
sample_indices = torch.tensor(
sample_indices, dtype=torch.int64, pin_memory=True
).to(pruned_states.device, non_blocking=True)
sample_indices, device=pruned_states.device, dtype=torch.int64
)
input_logprob_indices = torch.tensor(
input_logprob_indices, dtype=torch.int64, pin_memory=True
).to(pruned_states.device, non_blocking=True)
input_logprob_indices, device=pruned_states.device, dtype=torch.int64
)
return (
pruned_states,
@@ -590,24 +590,19 @@ class LogitsProcessor(nn.Module):
def _expand_metadata_for_logprobs(
self, logits_metadata: LogitsMetadata, device: torch.device
):
# Avoid implicit device sync inside repeat_interleave by providing output_size,
# which we can compute from CPU metadata.
total_pruned_len = sum(logits_metadata.extend_logprob_pruned_lens_cpu)
pruned_lens = torch.tensor(
logits_metadata.extend_logprob_pruned_lens_cpu,
pin_memory=True,
).to(device, non_blocking=True)
device=device,
)
if logits_metadata.temp_scaled_logprobs:
logits_metadata.temperature = torch.repeat_interleave(
logits_metadata.temperature.view(-1),
pruned_lens,
output_size=total_pruned_len,
).view(-1, 1)
if logits_metadata.top_p_normalized_logprobs:
logits_metadata.top_p = torch.repeat_interleave(
logits_metadata.top_p,
pruned_lens,
output_size=total_pruned_len,
)
def process_input_logprobs(self, input_logits, logits_metadata: LogitsMetadata):