From b58fa60a1f4375d3a0e668287263cc94f3d8eab8 Mon Sep 17 00:00:00 2001 From: Lucia Fang <116399278+luccafong@users.noreply.github.com> Date: Fri, 1 May 2026 12:52:46 -0700 Subject: [PATCH] [core/attention] Add SGLANG_FLASHINFER_USE_PAGED env to force paged wrapper (#24165) --- python/sglang/srt/environ.py | 1 + .../sglang/srt/layers/attention/flashinfer_backend.py | 10 +++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4c201d8f8..a4997b3cf 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -358,6 +358,7 @@ class Envs: # Flashinfer SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True) + SGLANG_FLASHINFER_USE_PAGED = EnvBool(False) # Default to the pick from flashinfer SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024) # Skip-softmax threshold scale factor for TRT-LLM attention (prefill and decode separately). diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 064e8adc7..27705a4b8 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -192,6 +192,8 @@ class FlashInferAttnBackend(AttentionBackend): self.disable_cuda_graph_kv_split = True envs.SGLANG_FLASHINFER_WORKSPACE_SIZE.set(2048 * 1024 * 1024) + self.use_paged = envs.SGLANG_FLASHINFER_USE_PAGED.get() + # Allocate buffers global global_workspace_buffer if global_workspace_buffer is None: @@ -489,7 +491,9 @@ class FlashInferAttnBackend(AttentionBackend): extend_no_prefix = False else: 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) @@ -698,7 +702,7 @@ class FlashInferAttnBackend(AttentionBackend): seq_lens_sum, prefix_lens=seq_lens - self.dllm_config.block_size, prefill_wrappers=prefill_wrappers, - use_ragged=True, + use_ragged=not self.use_paged, encoder_lens=encoder_lens, spec_info=None, ) @@ -762,7 +766,7 @@ class FlashInferAttnBackend(AttentionBackend): seq_lens_sum, prefix_lens=seq_lens - self.dllm_config.block_size, 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, spec_info=None, )