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

[core/attention] Add SGLANG_FLASHINFER_USE_PAGED env to force paged wrapper (#24165)

Browse Source
This commit is contained in:
Lucia Fang
2026-05-01 12:52:46 -07:00
committed by GitHub
parent c7566a860b
commit b58fa60a1f
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+1
View File
@@ -358,6 +358,7 @@ class Envs:
# Flashinfer # Flashinfer
SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True) SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True)
SGLANG_FLASHINFER_USE_PAGED = EnvBool(False)
# Default to the pick from flashinfer # Default to the pick from flashinfer
SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024)
# Skip-softmax threshold scale factor for TRT-LLM attention (prefill and decode separately). # Skip-softmax threshold scale factor for TRT-LLM attention (prefill and decode separately).
python/sglang/srt/layers/attention/flashinfer_backend.py
+7 -3
View File
@@ -192,6 +192,8 @@ class FlashInferAttnBackend(AttentionBackend):
self.disable_cuda_graph_kv_split = True self.disable_cuda_graph_kv_split = True
envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(2048 * 1024 * 1024) envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(2048 * 1024 * 1024)
self.use_paged = envs.SGLANG_FLASHINFER_USE_PAGED.get()
# Allocate buffers # Allocate buffers
global global_workspace_buffer global global_workspace_buffer
if global_workspace_buffer is None: if global_workspace_buffer is None:
@@ -489,7 +491,9 @@ class FlashInferAttnBackend(AttentionBackend):
extend_no_prefix = False extend_no_prefix = False
else: else:
use_ragged = ( use_ragged = (
not self.enable_deterministic and not is_in_piecewise_cuda_graph() not self.enable_deterministic
and not is_in_piecewise_cuda_graph()
and not self.use_paged
) )
extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu) extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
@@ -698,7 +702,7 @@ class FlashInferAttnBackend(AttentionBackend):
seq_lens_sum, seq_lens_sum,
prefix_lens=seq_lens - self.dllm_config.block_size, prefix_lens=seq_lens - self.dllm_config.block_size,
prefill_wrappers=prefill_wrappers, prefill_wrappers=prefill_wrappers,
use_ragged=True, use_ragged=not self.use_paged,
encoder_lens=encoder_lens, encoder_lens=encoder_lens,
spec_info=None, spec_info=None,
) )
@@ -762,7 +766,7 @@ class FlashInferAttnBackend(AttentionBackend):
seq_lens_sum, seq_lens_sum,
prefix_lens=seq_lens - self.dllm_config.block_size, prefix_lens=seq_lens - self.dllm_config.block_size,
prefill_wrappers=self.prefill_cuda_graph_metadata[bs], prefill_wrappers=self.prefill_cuda_graph_metadata[bs],
use_ragged=True, use_ragged=not self.use_paged,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None, encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=None, spec_info=None,
) )