sglang 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -15,6 +15,8 @@ limitations under the License.
 
 """Radix attention."""
 
+from typing import Optional
+
 import torch
 from flashinfer.cascade import merge_state
 from torch import nn
@@ -34,6 +36,7 @@ class RadixAttention(nn.Module):
         scaling: float,
         num_kv_heads: int,
         layer_id: int,
+        sliding_window_size: Optional[int] = None,
         logit_cap: int = -1,
         v_head_dim: int = -1,
     ):
@@ -46,6 +49,7 @@ class RadixAttention(nn.Module):
         self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim
         self.scaling = scaling
         self.layer_id = layer_id
+        self.sliding_window_size = sliding_window_size if sliding_window_size else -1
 
         if (
             not global_server_args_dict.get("disable_flashinfer", False)
@@ -113,14 +117,25 @@ class RadixAttention(nn.Module):
         return o
 
     def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
+        # using two wrappers is unnecessary in the current PR, but are prepared for future PRs
+        prefill_wrapper_paged = input_metadata.flashinfer_prefill_wrapper_paged
+        if self.sliding_window_size != -1:
+            prefill_wrapper_paged = prefill_wrapper_paged[0]
+        else:
+            if isinstance(prefill_wrapper_paged, list):
+                prefill_wrapper_paged = prefill_wrapper_paged[1]
+
         if not input_metadata.flashinfer_use_ragged:
-            self.store_kv_cache(k, v, input_metadata)
+            if k is not None:
+                assert v is not None
+                self.store_kv_cache(k, v, input_metadata)
 
-            o = input_metadata.flashinfer_prefill_wrapper_paged.forward(
+            o = prefill_wrapper_paged.forward(
                 q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
                 input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
                 causal=True,
                 sm_scale=self.scaling,
+                window_left=self.sliding_window_size,
                 logits_soft_cap=self.logit_cap,
             )
         else:
@@ -138,14 +153,12 @@ class RadixAttention(nn.Module):
             if input_metadata.extend_no_prefix:
                 o = o1
             else:
-                o2, s2 = (
-                    input_metadata.flashinfer_prefill_wrapper_paged.forward_return_lse(
-                        q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                        input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
-                        causal=False,
-                        sm_scale=self.scaling,
-                        logits_soft_cap=self.logit_cap,
-                    )
+                o2, s2 = prefill_wrapper_paged.forward_return_lse(
+                    q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
+                    input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
+                    causal=False,
+                    sm_scale=self.scaling,
+                    logits_soft_cap=self.logit_cap,
                 )
 
                 o, _ = merge_state(o1, s1, o2, s2)
@@ -158,9 +171,18 @@ class RadixAttention(nn.Module):
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        self.store_kv_cache(k, v, input_metadata)
+        decode_wrapper = input_metadata.flashinfer_decode_wrapper
+        if self.sliding_window_size != -1:
+            decode_wrapper = decode_wrapper[0]
+        else:
+            if isinstance(decode_wrapper, list):
+                decode_wrapper = decode_wrapper[1]
+
+        if k is not None:
+            assert v is not None
+            self.store_kv_cache(k, v, input_metadata)
 
-        o = input_metadata.flashinfer_decode_wrapper.forward(
+        o = decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
             sm_scale=self.scaling,
@@ -170,8 +192,10 @@ class RadixAttention(nn.Module):
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def forward(self, q, k, v, input_metadata: InputMetadata):
-        k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
-        v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
+        if k is not None:
+            assert v is not None
+            k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
+            v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
 
         if input_metadata.forward_mode == ForwardMode.EXTEND:
             return self.extend_forward(q, k, v, input_metadata)

sglang/srt/managers/schedule_batch.py

@@ -235,10 +235,12 @@ class Req:
             return
 
         last_token_id = self.output_ids[-1]
-        if self.tokenizer is None:
-            matched_eos = last_token_id in self.sampling_params.stop_token_ids
-        else:
-            matched_eos = last_token_id == self.tokenizer.eos_token_id
+
+        matched_eos = last_token_id in self.sampling_params.stop_token_ids
+
+        if self.tokenizer is not None:
+            matched_eos |= last_token_id == self.tokenizer.eos_token_id
+
         if matched_eos and not self.sampling_params.ignore_eos:
             self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id)
             return
@@ -383,7 +385,7 @@ class ScheduleBatch:
 
         return out_cache_loc
 
-    def batch_sampling_params(self, vocab_size, int_token_logit_bias):
+    def batch_sampling_params(self, vocab_size):
         device = "cuda"
         bs, reqs = self.batch_size(), self.reqs
         self.temperatures = torch.tensor(
@@ -419,15 +421,8 @@ class ScheduleBatch:
 
         # Handle logit bias but only allocate when needed
         self.logit_bias = None
-        for i in range(bs):
-            if reqs[i].sampling_params.dtype == "int":
-                if self.logit_bias is None:
-                    self.logit_bias = torch.zeros(
-                        (bs, vocab_size), dtype=torch.float32, device=device
-                    )
-                self.logit_bias[i][: len(int_token_logit_bias)] = int_token_logit_bias
 
-    def prepare_for_extend(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
+    def prepare_for_extend(self, vocab_size: int):
         bs = self.batch_size()
         reqs = self.reqs
         input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
@@ -466,7 +461,7 @@ class ScheduleBatch:
         self.out_cache_loc = out_cache_loc
         self.top_logprobs_nums = [r.top_logprobs_num for r in reqs]
 
-        self.batch_sampling_params(vocab_size, int_token_logit_bias)
+        self.batch_sampling_params(vocab_size)
 
     def check_decode_mem(self):
         bs = self.batch_size()

sglang/srt/managers/tokenizer_manager.py

@@ -507,7 +507,7 @@ class TokenizerManager:
             if obj.is_single:
                 self.abort_request(obj.rid)
             else:
-                for rid in obj.rids:
+                for rid in obj.rid:
                     self.abort_request(rid)
 
         background_tasks = BackgroundTasks()

sglang/srt/managers/tp_worker.py

@@ -54,7 +54,6 @@ from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
-    get_int_token_logit_bias,
     is_multimodal_model,
     set_random_seed,
     suppress_other_loggers,
@@ -132,9 +131,6 @@ class ModelTpServer:
             ),
             self.model_runner.req_to_token_pool.size - 1,
         )
-        self.int_token_logit_bias = torch.tensor(
-            get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
-        )
         self.max_req_input_len = min(
             self.model_config.context_len - 1,
             self.max_total_num_tokens - 1,
@@ -442,9 +438,7 @@ class ModelTpServer:
 
     def forward_prefill_batch(self, batch: ScheduleBatch):
         # Build batch tensors
-        batch.prepare_for_extend(
-            self.model_config.vocab_size, self.int_token_logit_bias
-        )
+        batch.prepare_for_extend(self.model_config.vocab_size)
 
         if self.model_runner.is_generation:
             # Forward and sample the next tokens

sglang/srt/model_executor/cuda_graph_runner.py

@@ -98,8 +98,8 @@ class CudaGraphRunner:
         self.req_pool_indices = torch.zeros(
             (self.max_bs,), dtype=torch.int32, device="cuda"
         )
-        self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.position_ids_offsets = torch.zeros(
+        self.seq_lens = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        self.position_ids_offsets = torch.ones(
             (self.max_bs,), dtype=torch.int32, device="cuda"
         )
         self.out_cache_loc = torch.zeros(
@@ -107,9 +107,6 @@ class CudaGraphRunner:
         )
 
         # FlashInfer inputs
-        self.flashinfer_workspace_buffer = (
-            self.model_runner.flashinfer_workspace_buffers[0]
-        )
         self.flashinfer_kv_indptr = torch.zeros(
             (self.max_bs + 1,), dtype=torch.int32, device="cuda"
         )
@@ -121,6 +118,23 @@ class CudaGraphRunner:
         self.flashinfer_kv_last_page_len = torch.ones(
             (self.max_bs,), dtype=torch.int32, device="cuda"
         )
+        if model_runner.sliding_window_size is None:
+            self.flashinfer_workspace_buffer = (
+                self.model_runner.flashinfer_workspace_buffer
+            )
+        else:
+            self.flashinfer_workspace_buffer = (
+                self.model_runner.flashinfer_workspace_buffer
+            )
+
+            self.flashinfer_kv_indptr = [
+                self.flashinfer_kv_indptr,
+                self.flashinfer_kv_indptr.clone(),
+            ]
+            self.flashinfer_kv_indices = [
+                self.flashinfer_kv_indices,
+                self.flashinfer_kv_indices.clone(),
+            ]
 
         self.compile_bs = [1, 2, 4, 8, 16, 24, 32] if use_torch_compile else []
 
@@ -171,15 +185,32 @@ class CudaGraphRunner:
             use_tensor_cores = True
         else:
             use_tensor_cores = False
-        flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffer,
-            "NHD",
-            use_cuda_graph=True,
-            use_tensor_cores=use_tensor_cores,
-            paged_kv_indptr_buffer=self.flashinfer_kv_indptr[: bs + 1],
-            paged_kv_indices_buffer=self.flashinfer_kv_indices,
-            paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[:bs],
-        )
+        if self.model_runner.sliding_window_size is None:
+            flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
+                self.flashinfer_workspace_buffer,
+                "NHD",
+                use_cuda_graph=True,
+                use_tensor_cores=use_tensor_cores,
+                paged_kv_indptr_buffer=self.flashinfer_kv_indptr[: bs + 1],
+                paged_kv_indices_buffer=self.flashinfer_kv_indices,
+                paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[:bs],
+            )
+        else:
+            flashinfer_decode_wrapper = []
+            for i in range(2):
+                flashinfer_decode_wrapper.append(
+                    BatchDecodeWithPagedKVCacheWrapper(
+                        self.flashinfer_workspace_buffer,
+                        "NHD",
+                        use_cuda_graph=True,
+                        use_tensor_cores=use_tensor_cores,
+                        paged_kv_indptr_buffer=self.flashinfer_kv_indptr[i][: bs + 1],
+                        paged_kv_indices_buffer=self.flashinfer_kv_indices[i],
+                        paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[
+                            :bs
+                        ],
+                    )
+                )
         update_flashinfer_indices(
             ForwardMode.DECODE,
             self.model_runner,
@@ -201,7 +232,7 @@ class CudaGraphRunner:
                 out_cache_loc=out_cache_loc,
                 return_logprob=False,
                 top_logprobs_nums=0,
-                positions=(seq_lens - 1).to(torch.int64),
+                positions=(seq_lens - 1 + position_ids_offsets).to(torch.int64),
                 flashinfer_decode_wrapper=flashinfer_decode_wrapper,
             )
 
@@ -225,8 +256,8 @@ class CudaGraphRunner:
         index = bisect.bisect_left(self.batch_size_list, raw_bs)
         bs = self.batch_size_list[index]
         if bs != raw_bs:
-            self.seq_lens.fill_(1)
-            self.position_ids_offsets.zero_()
+            self.seq_lens.zero_()
+            self.position_ids_offsets.fill_(1)
             self.out_cache_loc.zero_()
 
         # Common inputs

sglang/srt/model_executor/forward_batch_info.py

@@ -16,7 +16,7 @@ limitations under the License.
 """ModelRunner runs the forward passes of the models."""
 from dataclasses import dataclass
 from enum import IntEnum, auto
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
 import numpy as np
 import torch
@@ -194,6 +194,7 @@ class InputMetadata:
             if (
                 forward_mode != ForwardMode.DECODE
                 and int(torch.sum(ret.seq_lens)) > 4096
+                and model_runner.sliding_window_size is None
             ):
                 flashinfer_use_ragged = True
             ret.init_flashinfer_handlers(
@@ -216,7 +217,10 @@ class InputMetadata:
             self.triton_max_extend_len = int(torch.max(extend_seq_lens))
 
     def init_flashinfer_handlers(
-        self, model_runner, prefix_lens, flashinfer_use_ragged
+        self,
+        model_runner,
+        prefix_lens,
+        flashinfer_use_ragged,
     ):
         update_flashinfer_indices(
             self.forward_mode,
@@ -255,65 +259,135 @@ def update_flashinfer_indices(
     head_dim = model_runner.model_config.head_dim
     batch_size = len(req_pool_indices)
 
-    if flashinfer_use_ragged:
-        paged_kernel_lens = prefix_lens
-    else:
-        paged_kernel_lens = seq_lens
-
-    kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
-    kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
-    req_pool_indices_cpu = req_pool_indices.cpu().numpy()
-    paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
-    kv_indices = torch.cat(
-        [
-            model_runner.req_to_token_pool.req_to_token[
-                req_pool_indices_cpu[i], : paged_kernel_lens_cpu[i]
-            ]
-            for i in range(batch_size)
-        ],
-        dim=0,
-    ).contiguous()
-    kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
-
-    if forward_mode == ForwardMode.DECODE:
-        # CUDA graph uses different flashinfer_decode_wrapper
-        if flashinfer_decode_wrapper is None:
-            flashinfer_decode_wrapper = model_runner.flashinfer_decode_wrapper
-
-        flashinfer_decode_wrapper.end_forward()
-        flashinfer_decode_wrapper.begin_forward(
-            kv_indptr,
-            kv_indices,
-            kv_last_page_len,
-            num_qo_heads,
-            num_kv_heads,
-            head_dim,
-            1,
-        )
-    else:
-        # extend part
-        qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
-        qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
-
+    if model_runner.sliding_window_size is None:
         if flashinfer_use_ragged:
-            model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
-            model_runner.flashinfer_prefill_wrapper_ragged.begin_forward(
-                qo_indptr,
-                qo_indptr,
+            paged_kernel_lens = prefix_lens
+        else:
+            paged_kernel_lens = seq_lens
+
+        kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
+        kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+        req_pool_indices_cpu = req_pool_indices.cpu().numpy()
+        paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
+        kv_indices = torch.cat(
+            [
+                model_runner.req_to_token_pool.req_to_token[
+                    req_pool_indices_cpu[i], : paged_kernel_lens_cpu[i]
+                ]
+                for i in range(batch_size)
+            ],
+            dim=0,
+        ).contiguous()
+        kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
+
+        if forward_mode == ForwardMode.DECODE:
+            # CUDA graph uses different flashinfer_decode_wrapper
+            if flashinfer_decode_wrapper is None:
+                flashinfer_decode_wrapper = model_runner.flashinfer_decode_wrapper
+
+            flashinfer_decode_wrapper.end_forward()
+            flashinfer_decode_wrapper.begin_forward(
+                kv_indptr,
+                kv_indices,
+                kv_last_page_len,
                 num_qo_heads,
                 num_kv_heads,
                 head_dim,
+                1,
             )
+        else:
+            # extend part
+            qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
+            qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+
+            if flashinfer_use_ragged:
+                model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
+                model_runner.flashinfer_prefill_wrapper_ragged.begin_forward(
+                    qo_indptr,
+                    qo_indptr,
+                    num_qo_heads,
+                    num_kv_heads,
+                    head_dim,
+                )
 
-        # cached part
-        model_runner.flashinfer_prefill_wrapper_paged.end_forward()
-        model_runner.flashinfer_prefill_wrapper_paged.begin_forward(
-            qo_indptr,
-            kv_indptr,
-            kv_indices,
-            kv_last_page_len,
-            num_qo_heads,
-            num_kv_heads,
-            head_dim,
-            1,
-        )
+            # cached part
+            model_runner.flashinfer_prefill_wrapper_paged.end_forward()
+            model_runner.flashinfer_prefill_wrapper_paged.begin_forward(
+                qo_indptr,
+                kv_indptr,
+                kv_indices,
+                kv_last_page_len,
+                num_qo_heads,
+                num_kv_heads,
+                head_dim,
+                1,
+            )
+    else:
+        # window attention use paged only
+        kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
+        for wrapper_id in range(2):
+            if wrapper_id == 0:
+                if forward_mode == ForwardMode.DECODE:
+                    paged_kernel_lens = torch.minimum(
+                        seq_lens, torch.tensor(model_runner.sliding_window_size + 1)
+                    )
+                else:
+                    paged_kernel_lens = torch.minimum(
+                        seq_lens,
+                        torch.tensor(model_runner.sliding_window_size)
+                        + seq_lens
+                        - prefix_lens,
+                    )
+            else:
+                paged_kernel_lens = seq_lens
+
+            kv_start_idx = seq_lens - paged_kernel_lens
+
+            kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
+            kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+            req_pool_indices_cpu = req_pool_indices.cpu().numpy()
+            paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
+            kv_indices = torch.cat(
+                [
+                    model_runner.req_to_token_pool.req_to_token[
+                        req_pool_indices_cpu[i],
+                        kv_start_idx[i] : kv_start_idx[i] + paged_kernel_lens_cpu[i],
+                    ]
+                    for i in range(batch_size)
+                ],
+                dim=0,
+            ).contiguous()
+
+            if forward_mode == ForwardMode.DECODE:
+                # CUDA graph uses different flashinfer_decode_wrapper
+                if flashinfer_decode_wrapper is None:
+                    flashinfer_decode_wrapper = model_runner.flashinfer_decode_wrapper
+
+                flashinfer_decode_wrapper[wrapper_id].end_forward()
+                flashinfer_decode_wrapper[wrapper_id].begin_forward(
+                    kv_indptr,
+                    kv_indices,
+                    kv_last_page_len,
+                    num_qo_heads,
+                    num_kv_heads,
+                    head_dim,
+                    1,
+                )
+            else:
+                # extend part
+                qo_indptr = torch.zeros(
+                    (batch_size + 1,), dtype=torch.int32, device="cuda"
+                )
+                qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
+
+                model_runner.flashinfer_prefill_wrapper_paged[wrapper_id].end_forward()
+                model_runner.flashinfer_prefill_wrapper_paged[wrapper_id].begin_forward(
+                    qo_indptr,
+                    kv_indptr,
+                    kv_indices,
+                    kv_last_page_len,
+                    num_qo_heads,
+                    num_kv_heads,
+                    head_dim,
+                    1,
+                )

sglang/srt/model_executor/model_runner.py

@@ -38,6 +38,7 @@ from vllm.distributed import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm.model_executor.model_loader import get_model
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
@@ -53,7 +54,7 @@ from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
     get_available_gpu_memory,
     is_generation_model,
-    is_llama3_405b_fp8,
+    is_llama3_405b_fp8_head_16,
     is_multimodal_model,
     monkey_patch_vllm_dummy_weight_loader,
     monkey_patch_vllm_p2p_access_check,
@@ -158,7 +159,7 @@ class ModelRunner:
             skip_tokenizer_init=True,
         )
 
-        if is_llama3_405b_fp8(self.model_config) and self.tp_size <= 8:
+        if is_llama3_405b_fp8_head_16(self.model_config) and self.tp_size <= 8:
             # A temporary hack to fix the num_heads for meta-llama/Meta-Llama-3.1-405B-FP8 checkpoints
             self.model_config.hf_config.num_key_value_heads = 8
             vllm_model_config.hf_config.num_key_value_heads = 8
@@ -168,15 +169,6 @@ class ModelRunner:
         if self.model_config.model_overide_args is not None:
             vllm_model_config.hf_config.update(self.model_config.model_overide_args)
 
-        if (
-            self.server_args.efficient_weight_load
-            and "llama" in self.server_args.model_path.lower()
-            and self.server_args.quantization == "fp8"
-        ):
-            from sglang.srt.model_loader.model_loader import get_model
-        else:
-            from vllm.model_executor.model_loader import get_model
-
         self.model = get_model(
             model_config=vllm_model_config,
             device_config=device_config,
@@ -187,6 +179,11 @@ class ModelRunner:
             scheduler_config=None,
             cache_config=None,
         )
+        self.sliding_window_size = (
+            self.model.get_window_size()
+            if hasattr(self.model, "get_window_size")
+            else None
+        )
         self.is_generation = is_generation_model(
             self.model_config.hf_config.architectures
         )
@@ -296,6 +293,9 @@ class ModelRunner:
 
     def init_flashinfer(self):
         if self.server_args.disable_flashinfer:
+            assert (
+                self.sliding_window_size is None
+            ), "turn on flashinfer to support window attention"
             self.flashinfer_prefill_wrapper_ragged = None
             self.flashinfer_prefill_wrapper_paged = None
             self.flashinfer_decode_wrapper = None
@@ -309,20 +309,47 @@ class ModelRunner:
         else:
             use_tensor_cores = False
 
-        self.flashinfer_workspace_buffers = torch.empty(
-            2, global_config.flashinfer_workspace_size, dtype=torch.uint8, device="cuda"
-        )
-        self.flashinfer_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0], "NHD"
-        )
-        self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[1], "NHD"
-        )
-        self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffers[0],
-            "NHD",
-            use_tensor_cores=use_tensor_cores,
-        )
+        if self.sliding_window_size is None:
+            self.flashinfer_workspace_buffer = torch.empty(
+                global_config.flashinfer_workspace_size,
+                dtype=torch.uint8,
+                device="cuda",
+            )
+            self.flashinfer_prefill_wrapper_ragged = (
+                BatchPrefillWithRaggedKVCacheWrapper(
+                    self.flashinfer_workspace_buffer, "NHD"
+                )
+            )
+            self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
+                self.flashinfer_workspace_buffer, "NHD"
+            )
+            self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
+                self.flashinfer_workspace_buffer,
+                "NHD",
+                use_tensor_cores=use_tensor_cores,
+            )
+        else:
+            self.flashinfer_workspace_buffer = torch.empty(
+                global_config.flashinfer_workspace_size,
+                dtype=torch.uint8,
+                device="cuda",
+            )
+            self.flashinfer_prefill_wrapper_ragged = None
+            self.flashinfer_prefill_wrapper_paged = []
+            self.flashinfer_decode_wrapper = []
+            for i in range(2):
+                self.flashinfer_prefill_wrapper_paged.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
+                        self.flashinfer_workspace_buffer, "NHD"
+                    )
+                )
+                self.flashinfer_decode_wrapper.append(
+                    BatchDecodeWithPagedKVCacheWrapper(
+                        self.flashinfer_workspace_buffer,
+                        "NHD",
+                        use_tensor_cores=use_tensor_cores,
+                    )
+                )
 
     def init_cuda_graphs(self):
         from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
@@ -358,7 +385,9 @@ class ModelRunner:
             return self.cuda_graph_runner.replay(batch)
 
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, ForwardMode.DECODE
+            self,
+            batch,
+            ForwardMode.DECODE,
         )
 
         return self.model.forward(
@@ -368,7 +397,9 @@ class ModelRunner:
     @torch.inference_mode()
     def forward_extend(self, batch: ScheduleBatch):
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, forward_mode=ForwardMode.EXTEND
+            self,
+            batch,
+            forward_mode=ForwardMode.EXTEND,
         )
         return self.model.forward(
             batch.input_ids, input_metadata.positions, input_metadata
@@ -377,7 +408,9 @@ class ModelRunner:
     @torch.inference_mode()
     def forward_extend_multi_modal(self, batch: ScheduleBatch):
         input_metadata = InputMetadata.from_schedule_batch(
-            self, batch, forward_mode=ForwardMode.EXTEND
+            self,
+            batch,
+            forward_mode=ForwardMode.EXTEND,
         )
         return self.model.forward(
             batch.input_ids,