EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Adjust default mem fraction to avoid OOM (#823)

Browse Source
This commit is contained in:
Ying Sheng
2024-07-30 01:58:31 -07:00
committed by GitHub
parent ae5c0fc442
commit e7487b08bc
4 changed files with 22 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -781,7 +781,7 @@ class InputMetadata:
flashinfer_prefill_wrapper_ragged: "BatchPrefillWithRaggedKVCacheWrapper" = None
flashinfer_prefill_wrapper_paged: "BatchPrefillWithPagedKVCacheWrapper" = None
flashinfer_decode_wrapper: "BatchDecodeWithPagedKVCacheWrapper" = None
use_ragged: bool = False
flashinfer_use_ragged: bool = False
@classmethod
def create(
@@ -797,10 +797,10 @@ class InputMetadata:
return_logprob=False,
skip_flashinfer_init=False,
):
use_ragged = False
flashinfer_use_ragged = False
if not skip_flashinfer_init and not model_runner.server_args.disable_flashinfer:
if forward_mode != ForwardMode.DECODE and int(torch.sum(seq_lens)) > 4096:
use_ragged = True
flashinfer_use_ragged = True
init_flashinfer_args(
forward_mode,
model_runner,
@@ -808,7 +808,7 @@ class InputMetadata:
seq_lens,
prefix_lens,
model_runner.flashinfer_decode_wrapper,
use_ragged,
flashinfer_use_ragged,
)
batch_size = len(req_pool_indices)
@@ -863,7 +863,7 @@ class InputMetadata:
flashinfer_prefill_wrapper_ragged=model_runner.flashinfer_prefill_wrapper_ragged,
flashinfer_prefill_wrapper_paged=model_runner.flashinfer_prefill_wrapper_paged,
flashinfer_decode_wrapper=model_runner.flashinfer_decode_wrapper,
use_ragged=use_ragged,
flashinfer_use_ragged=flashinfer_use_ragged,
)
if model_runner.server_args.disable_flashinfer:
@@ -884,7 +884,7 @@ def init_flashinfer_args(
seq_lens,
prefix_lens,
flashinfer_decode_wrapper,
use_ragged=False,
flashinfer_use_ragged=False,
):
"""Init auxiliary variables for FlashInfer attention backend."""
num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
@@ -893,7 +893,7 @@ def init_flashinfer_args(
batch_size = len(req_pool_indices)
total_num_tokens = int(torch.sum(seq_lens))
if use_ragged:
if flashinfer_use_ragged:
paged_kernel_lens = prefix_lens
else:
paged_kernel_lens = seq_lens
@@ -929,7 +929,7 @@ def init_flashinfer_args(
qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
if use_ragged:
if flashinfer_use_ragged:
model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
model_runner.flashinfer_prefill_wrapper_ragged.begin_forward(
qo_indptr,