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

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

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
This commit is contained in:
Leon Gao
2026-02-27 18:11:05 -08:00
committed by GitHub
parent 8c0f2d40bd
commit b5a8e4179e
7 changed files with 87 additions and 42 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, device=pruned_states.device, dtype=torch.int64
)
sample_indices, dtype=torch.int64, pin_memory=True
).to(pruned_states.device, non_blocking=True)
input_logprob_indices = torch.tensor(
input_logprob_indices, device=pruned_states.device, dtype=torch.int64
)
input_logprob_indices, dtype=torch.int64, pin_memory=True
).to(pruned_states.device, non_blocking=True)
return (
pruned_states,
@@ -590,19 +590,24 @@ 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,
device=device,
)
pin_memory=True,
).to(device, non_blocking=True)
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):