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

1
python/sglang/srt/batch_overlap/two_batch_overlap.py
View File

@@ -685,7 +685,6 @@ class TboForwardBatchPreparer:
for key in [
"forward_mode",
"is_extend_in_batch",
"all_extend_in_batch",
"return_logprob",
"req_to_token_pool",
"token_to_kv_pool",

2
python/sglang/srt/connector/__init__.py
View File

@@ -22,7 +22,7 @@ class ConnectorType(str, enum.Enum):
INSTANCE = "instance"
def create_remote_connector(url, device=None, **kwargs) -> BaseConnector:
def create_remote_connector(url, device, **kwargs) -> BaseConnector:
connector_type = parse_connector_type(url)
if connector_type == "redis":
return RedisConnector(url)

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):

36
python/sglang/srt/managers/schedule_batch.py
View File

@@ -1226,7 +1226,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
global_num_tokens: Optional[List[int]] = None
global_num_tokens_for_logprob: Optional[List[int]] = None
is_extend_in_batch: bool = False
all_extend_in_batch: bool = False
can_run_dp_cuda_graph: bool = False
tbo_split_seq_index: Optional[int] = None
global_forward_mode: Optional[ForwardMode] = None
@@ -1986,34 +1985,22 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self.seq_lens_sum += bs
if get_global_server_args().enable_mamba_extra_buffer():
# Build indices fully on GPU without scalar extraction.
# Each slice is shape [1]; cat -> [bs].
if len(self.reqs) == 0:
self.mamba_track_indices = torch.empty(
(0,), dtype=torch.int64, device=self.device
)
else:
self.mamba_track_indices = torch.cat(
[
(
req.mamba_ping_pong_track_buffer[1:]
if req.mamba_next_track_idx == 1
else req.mamba_ping_pong_track_buffer[:1]
)
for req in self.reqs
],
dim=0,
).to(torch.int64)
# Keep mask construction in the pinned-tensor form.
self.mamba_track_indices = torch.tensor(
[
req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx]
for req in self.reqs
],
dtype=torch.int64,
device=self.device,
)
self.mamba_track_mask = torch.tensor(
[
sl % get_global_server_args().mamba_track_interval == 0
for sl in self.seq_lens_cpu
],
dtype=torch.bool,
pin_memory=True,
).to(device=self.device, non_blocking=True)
device=self.device,
)
def maybe_wait_verify_done(self):
if self.is_spec_v2:
@@ -2183,7 +2170,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
global_num_tokens=self.global_num_tokens,
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
is_extend_in_batch=self.is_extend_in_batch,
all_extend_in_batch=self.all_extend_in_batch,
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
tbo_split_seq_index=self.tbo_split_seq_index,
global_forward_mode=self.global_forward_mode,
@@ -2241,7 +2227,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
global_num_tokens=self.global_num_tokens,
global_num_tokens_for_logprob=self.global_num_tokens_for_logprob,
can_run_dp_cuda_graph=self.can_run_dp_cuda_graph,
all_extend_in_batch=self.all_extend_in_batch,
is_extend_in_batch=self.is_extend_in_batch,
is_prefill_only=self.is_prefill_only,
seq_lens_cpu=self.seq_lens_cpu,
@@ -2346,7 +2331,6 @@ class ModelWorkerBatch:
global_num_tokens: Optional[List[int]]
global_num_tokens_for_logprob: Optional[List[int]]
is_extend_in_batch: bool
all_extend_in_batch: bool
can_run_dp_cuda_graph: bool
tbo_split_seq_index: Optional[int]
global_forward_mode: Optional[ForwardMode]

69
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -343,18 +343,10 @@ class MambaPool:
select_index = self.free_slots[:need_size]
self.free_slots = self.free_slots[need_size:]
# clear at alloc time — expand a scalar GPU zero to the right shape, no CPU-GPU sync
# clear at alloc time, fill allocated slots with zeros
for i in range(len(self.mamba_cache.conv)):
t = self.mamba_cache.conv[i]
z = torch.zeros(1, dtype=t.dtype, device=t.device).expand(
t.shape[0], need_size, *t.shape[2:]
)
t[:, select_index] = z
t = self.mamba_cache.temporal
z = torch.zeros(1, dtype=t.dtype, device=t.device).expand(
t.shape[0], need_size, *t.shape[2:]
)
t[:, select_index] = z
self.mamba_cache.conv[i][:, select_index] = 0
self.mamba_cache.temporal[:, select_index] = 0
return select_index
@@ -522,8 +514,8 @@ class HybridReqToTokenPool(ReqToTokenPool):
if select_index is None:
return None
mamba_indices: list[torch.Tensor] = []
mamba_ping_pong_track_buffers: list[torch.Tensor] = []
mamba_index = []
mamba_ping_pong_track_buffer_list = []
for req in reqs:
mid = None
if req.mamba_pool_idx is not None: # for radix cache
@@ -535,7 +527,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_pool.available_size()=}, {len(reqs)=}"
mid = mid[0]
req.mamba_pool_idx = mid
mamba_indices.append(mid)
mamba_index.append(mid)
if self.enable_mamba_extra_buffer:
if req.mamba_ping_pong_track_buffer is None:
req.mamba_ping_pong_track_buffer = self.mamba_pool.alloc(
@@ -545,22 +537,26 @@ class HybridReqToTokenPool(ReqToTokenPool):
req.mamba_ping_pong_track_buffer is not None
), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
req.mamba_next_track_idx = 0
mamba_ping_pong_track_buffers.append(req.mamba_ping_pong_track_buffer)
mamba_ping_pong_track_buffer_list.append(
req.mamba_ping_pong_track_buffer.tolist()
)
assert len(select_index) == len(
mamba_indices
mamba_index
), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size."
if self.enable_mamba_extra_buffer:
assert len(select_index) == len(
mamba_ping_pong_track_buffers
mamba_ping_pong_track_buffer_list
), f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio."
mamba_index_tensor = torch.stack(mamba_indices).to(dtype=torch.int32)
self.req_index_to_mamba_index_mapping[select_index] = mamba_index_tensor
self.req_index_to_mamba_index_mapping[select_index] = torch.tensor(
mamba_index, dtype=torch.int32, device=self.device
)
if self.enable_mamba_extra_buffer:
ping_pong_tensor = torch.stack(mamba_ping_pong_track_buffers).to(
dtype=torch.int32
)
self.req_index_to_mamba_ping_pong_track_buffer_mapping[select_index] = (
ping_pong_tensor
torch.tensor(
mamba_ping_pong_track_buffer_list,
dtype=torch.int32,
device=self.device,
)
)
return select_index
@@ -597,28 +593,11 @@ class HybridReqToTokenPool(ReqToTokenPool):
0,
1,
], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}"
# Avoid Python-list advanced indexing on a device tensor.
# The ping-pong buffer size is either 2 (normal) or 1 (spec decode).
if self.mamba_ping_pong_track_buffer_size == 2:
idx_to_free = 1 - mamba_ping_pong_track_buffer_to_keep
mamba_ping_pong_track_buffer_to_free = (
mamba_ping_pong_track_buffer_to_free[
idx_to_free : idx_to_free + 1
]
)
else:
assert self.mamba_ping_pong_track_buffer_size == 1, (
f"Unexpected mamba_ping_pong_track_buffer_size="
f"{self.mamba_ping_pong_track_buffer_size}"
)
assert mamba_ping_pong_track_buffer_to_keep == 0, (
"mamba_ping_pong_track_buffer_to_keep must be 0 when "
"mamba_ping_pong_track_buffer_size is 1"
)
# Keep the only slot, so free nothing.
mamba_ping_pong_track_buffer_to_free = (
mamba_ping_pong_track_buffer_to_free[0:0]
)
idx_to_free = list(range(self.mamba_ping_pong_track_buffer_size))
idx_to_free.remove(mamba_ping_pong_track_buffer_to_keep)
mamba_ping_pong_track_buffer_to_free = (
mamba_ping_pong_track_buffer_to_free[idx_to_free]
)
self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free)
def clear(self):

2
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -338,7 +338,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
dp_local_num_tokens: Optional[torch.Tensor] = None # cached info at runtime
global_dp_buffer_len: Optional[int] = None
is_extend_in_batch: bool = False
all_extend_in_batch: bool = False
can_run_dp_cuda_graph: bool = False
global_forward_mode: Optional[ForwardMode] = None
@@ -405,7 +404,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
top_logprobs_nums=batch.top_logprobs_nums,
token_ids_logprobs=batch.token_ids_logprobs,
is_extend_in_batch=batch.is_extend_in_batch,
all_extend_in_batch=batch.all_extend_in_batch,
can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph,
global_forward_mode=batch.global_forward_mode,
is_prefill_only=batch.is_prefill_only,

2
python/sglang/srt/model_executor/model_runner.py
View File

@@ -1134,7 +1134,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
"""Update engine weights in-place from the disk."""
logger.info(
f"Update engine weights online from disk begin. "
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id, empty_cache=False):.2f} GB"
f"avail mem={get_available_gpu_memory(self.device, self.gpu_id):.2f} GB"
)
target_device = torch.device(self.device)