sglang 0.1.22__py3-none-any.whl → 0.1.25__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -85,32 +85,47 @@ class RadixAttention(nn.Module):
         return o
 
     def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        o1, s1 = input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
-            q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
-            v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
-            causal=True,
-            sm_scale=self.scaling,
-            logits_soft_cap=self.logit_cap,
-        )
+        if not input_metadata.use_ragged:
+            self.store_kv_cache(k, v, input_metadata)
 
-        if input_metadata.extend_no_prefix:
-            o = o1
-        else:
-            o2, s2 = input_metadata.flashinfer_prefill_wrapper_paged.forward_return_lse(
+            o = input_metadata.flashinfer_prefill_wrapper_paged.forward(
                 q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                input_metadata.token_to_kv_pool.kv_data[self.layer_id],
-                causal=False,
+                input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
+                causal=True,
                 sm_scale=self.scaling,
                 logits_soft_cap=self.logit_cap,
             )
+        else:
+            o1, s1 = (
+                input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
+                    q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
+                    k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
+                    v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
+                    causal=True,
+                    sm_scale=self.scaling,
+                    logits_soft_cap=self.logit_cap,
+                )
+            )
 
-            o, _ = merge_state(o1, s1, o2, s2)
+            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,
+                    )
+                )
 
-        self.store_kv_cache(k, v, input_metadata)
+                o, _ = merge_state(o1, s1, o2, s2)
+
+            self.store_kv_cache(k, v, input_metadata)
 
-        if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
-            torch.cuda.synchronize()
+            if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
+                torch.cuda.synchronize()
 
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
@@ -119,7 +134,7 @@ class RadixAttention(nn.Module):
 
         o = input_metadata.flashinfer_decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
             sm_scale=self.scaling,
             logits_soft_cap=self.logit_cap,
         )
@@ -136,33 +151,7 @@ class RadixAttention(nn.Module):
             return self.decode_forward(q, k, v, input_metadata)
 
     def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
-        kv_cache = input_metadata.token_to_kv_pool.kv_data[self.layer_id]
-        _store_kv_cache(cache_k, cache_v, kv_cache, input_metadata.out_cache_loc)
-
-
-try:
-
-    @torch.library.custom_op("mylib::store_kv_cache", mutates_args={"kv_cache"})
-    def _store_kv_cache(
-        k: torch.Tensor,
-        v: torch.Tensor,
-        kv_cache: torch.Tensor,
-        cache_loc: torch.Tensor,
-    ) -> None:
-        kv_cache[cache_loc, 0] = k
-        kv_cache[cache_loc, 1] = v
-
-    @_store_kv_cache.register_fake
-    def _(k, v, kv_cache, cache_loc):
-        pass
-
-except:
-
-    def _store_kv_cache(
-        k: torch.Tensor,
-        v: torch.Tensor,
-        kv_cache: torch.Tensor,
-        cache_loc: torch.Tensor,
-    ) -> None:
-        kv_cache[cache_loc, 0] = k
-        kv_cache[cache_loc, 1] = v
+        k_cache = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
+        v_cache = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
+        k_cache[input_metadata.out_cache_loc] = cache_k
+        v_cache[input_metadata.out_cache_loc] = cache_v

sglang/srt/managers/controller/cuda_graph_runner.py

@@ -1,11 +1,13 @@
 """Run the model with cuda graph."""
 
 import bisect
+from contextlib import contextmanager
 
 import torch
 from flashinfer import BatchDecodeWithPagedKVCacheWrapper
 from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 from vllm.distributed.parallel_state import graph_capture
+from vllm.model_executor.custom_op import CustomOp
 
 from sglang.srt.layers.logits_processor import LogitProcessorOutput
 from sglang.srt.managers.controller.infer_batch import (
@@ -14,10 +16,44 @@ from sglang.srt.managers.controller.infer_batch import (
     InputMetadata,
     init_flashinfer_args,
 )
+from sglang.srt.utils import monkey_patch_vllm_all_gather
+
+
+def _to_torch(model: torch.nn.Module, reverse: bool = False):
+    for sub in model._modules.values():
+        if isinstance(sub, CustomOp):
+            if reverse:
+                sub._forward_method = sub.forward_cuda
+            else:
+                sub._forward_method = sub.forward_native
+        if isinstance(sub, torch.nn.Module):
+            _to_torch(sub, reverse)
+
+
+@contextmanager
+def patch_model(
+    model: torch.nn.Module, use_compile: bool, tp_group: "GroupCoordinator"
+):
+    backup_ca_comm = None
+
+    try:
+        if use_compile:
+            _to_torch(model)
+            monkey_patch_vllm_all_gather()
+            backup_ca_comm = tp_group.ca_comm
+            tp_group.ca_comm = None
+            yield torch.compile(model.forward, mode="max-autotune-no-cudagraphs")
+        else:
+            yield model.forward
+    finally:
+        if use_compile:
+            _to_torch(model, reverse=True)
+            monkey_patch_vllm_all_gather(reverse=True)
+            tp_group.ca_comm = backup_ca_comm
 
 
 class CudaGraphRunner:
-    def __init__(self, model_runner, max_batch_size_to_capture):
+    def __init__(self, model_runner, max_batch_size_to_capture, use_torch_compile):
         self.model_runner = model_runner
         self.graphs = {}
         self.input_buffers = {}
@@ -55,6 +91,8 @@ class CudaGraphRunner:
             (self.max_bs,), dtype=torch.int32, device="cuda"
         )
 
+        self.compile_bs = [1, 2, 4, 8, 16, 24, 32] if use_torch_compile else []
+
     def can_run(self, batch_size):
         return batch_size < self.max_bs
 
@@ -63,18 +101,23 @@ class CudaGraphRunner:
         with graph_capture() as graph_capture_context:
             self.stream = graph_capture_context.stream
             for bs in batch_size_list:
-                (
-                    graph,
-                    input_buffers,
-                    output_buffers,
-                    flashinfer_handler,
-                ) = self.capture_one_batch_size(bs)
-                self.graphs[bs] = graph
-                self.input_buffers[bs] = input_buffers
-                self.output_buffers[bs] = output_buffers
-                self.flashinfer_handlers[bs] = flashinfer_handler
-
-    def capture_one_batch_size(self, bs):
+                with patch_model(
+                    self.model_runner.model,
+                    bs in self.compile_bs,
+                    self.model_runner.tp_group,
+                ) as forward:
+                    (
+                        graph,
+                        input_buffers,
+                        output_buffers,
+                        flashinfer_handler,
+                    ) = self.capture_one_batch_size(bs, forward)
+                    self.graphs[bs] = graph
+                    self.input_buffers[bs] = input_buffers
+                    self.output_buffers[bs] = output_buffers
+                    self.flashinfer_handlers[bs] = flashinfer_handler
+
+    def capture_one_batch_size(self, bs, forward):
         graph = torch.cuda.CUDAGraph()
         stream = self.stream
 
@@ -127,9 +170,8 @@ class CudaGraphRunner:
                 skip_flashinfer_init=True,
             )
             input_metadata.flashinfer_decode_wrapper = flashinfer_decode_wrapper
-            return self.model_runner.model.forward(
-                input_ids, input_metadata.positions, input_metadata
-            )
+
+            return forward(input_ids, input_metadata.positions, input_metadata)
 
         for _ in range(2):
             run_once()

sglang/srt/managers/controller/infer_batch.py

@@ -9,6 +9,7 @@ import numpy as np
 import torch
 from flashinfer.sampling import top_k_top_p_sampling_from_probs
 
+from sglang.global_config import global_config
 from sglang.srt.constrained import RegexGuide
 from sglang.srt.constrained.jump_forward import JumpForwardMap
 from sglang.srt.managers.controller.radix_cache import RadixCache
@@ -431,7 +432,8 @@ class Batch:
 
     def retract_decode(self):
         sorted_indices = [i for i in range(len(self.reqs))]
-        # TODO(lsyin): improve the priority of retraction
+
+        # TODO(lsyin): improve retraction policy for radix cache
         sorted_indices.sort(
             key=lambda i: (
                 len(self.reqs[i].output_ids),
@@ -443,7 +445,17 @@ class Batch:
         retracted_reqs = []
         seq_lens_cpu = self.seq_lens.cpu().numpy()
         req_pool_indices_cpu = self.req_pool_indices.cpu().numpy()
-        while self.token_to_kv_pool.available_size() < len(self.reqs):
+        while (
+            self.token_to_kv_pool.available_size()
+            < len(sorted_indices) * global_config.retract_decode_steps
+        ):
+            if len(sorted_indices) == 1:
+                # Corner case: only one request left
+                assert (
+                    self.token_to_kv_pool.available_size() > 0
+                ), "No space left for only one request"
+                break
+
             idx = sorted_indices.pop()
             req = self.reqs[idx]
             retracted_reqs.append(req)
@@ -468,7 +480,16 @@ class Batch:
 
         self.filter_batch(sorted_indices)
 
-        return retracted_reqs
+        # Reqs in batch are filtered
+        total_decoded_tokens = sum(len(r.output_ids) for r in self.reqs)
+        total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in self.reqs)
+
+        new_estimate_ratio = (
+            total_decoded_tokens + global_config.retract_decode_steps * len(self.reqs)
+        ) / total_max_new_tokens
+        new_estimate_ratio = min(1.0, new_estimate_ratio)
+
+        return retracted_reqs, new_estimate_ratio
 
     def check_for_jump_forward(self, model_runner):
         jump_forward_reqs = []
@@ -668,18 +689,17 @@ class Batch:
 
         max_top_k_round, batch_size = 32, probs.shape[0]
         uniform_samples = torch.rand((max_top_k_round, batch_size), device=probs.device)
-        batch_next_token_ids, _ = top_k_top_p_sampling_from_probs(
+        batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
             probs, uniform_samples, self.top_ks, self.top_ps
         )
 
-        # FIXME: this is a temporary fix for the illegal token ids
-        illegal_mask = torch.logical_or(
-            batch_next_token_ids < 0, batch_next_token_ids >= probs.shape[-1]
-        )
-        if torch.any(illegal_mask):
-            warnings.warn("Illegal sampled token ids")
+        if torch.any(~success):
+            warnings.warn("Sampling failed, fallback to top_k=1 strategy")
             probs = probs.masked_fill(torch.isnan(probs), 0.0)
-            batch_next_token_ids = torch.argmax(probs, dim=-1)
+            argmax_ids = torch.argmax(probs, dim=-1)
+            batch_next_token_ids = torch.where(
+                success, batch_next_token_ids, argmax_ids
+            )
 
         if has_regex:
             batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy()
@@ -727,6 +747,7 @@ class InputMetadata:
     flashinfer_prefill_wrapper_ragged: "BatchPrefillWithRaggedKVCacheWrapper" = None
     flashinfer_prefill_wrapper_paged: "BatchPrefillWithPagedKVCacheWrapper" = None
     flashinfer_decode_wrapper: "BatchDecodeWithPagedKVCacheWrapper" = None
+    use_ragged: bool = False
 
     @classmethod
     def create(
@@ -742,7 +763,10 @@ class InputMetadata:
         return_logprob=False,
         skip_flashinfer_init=False,
     ):
+        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
             init_flashinfer_args(
                 forward_mode,
                 model_runner,
@@ -750,6 +774,7 @@ class InputMetadata:
                 seq_lens,
                 prefix_lens,
                 model_runner.flashinfer_decode_wrapper,
+                use_ragged,
             )
 
         batch_size = len(req_pool_indices)
@@ -804,6 +829,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,
         )
 
         if model_runner.server_args.disable_flashinfer:
@@ -824,17 +850,19 @@ def init_flashinfer_args(
     seq_lens,
     prefix_lens,
     flashinfer_decode_wrapper,
+    use_ragged=False,
 ):
     """Init auxiliary variables for FlashInfer attention backend."""
     num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
     num_kv_heads = model_runner.model_config.get_num_kv_heads(model_runner.tp_size)
     head_dim = model_runner.model_config.head_dim
     batch_size = len(req_pool_indices)
+    total_num_tokens = int(torch.sum(seq_lens))
 
-    if forward_mode == ForwardMode.DECODE:
-        paged_kernel_lens = seq_lens
-    else:
+    if 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)
@@ -867,14 +895,15 @@ 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)
 
-        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,
-        )
+        if 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()

sglang/srt/managers/controller/model_runner.py

@@ -15,6 +15,7 @@ from flashinfer import (
     BatchPrefillWithRaggedKVCacheWrapper,
 )
 from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+from torch.nn.parameter import Parameter
 from vllm.config import DeviceConfig, LoadConfig
 from vllm.config import ModelConfig as VllmModelConfig
 from vllm.distributed import (
@@ -22,7 +23,7 @@ from vllm.distributed import (
     init_distributed_environment,
     initialize_model_parallel,
 )
-from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
@@ -39,6 +40,18 @@ from sglang.srt.utils import (
 logger = logging.getLogger("srt.model_runner")
 
 
+def is_llama3_405b_fp8(model_config):
+    if (
+        model_config.hf_config.architectures[0] == "LlamaForCausalLM"
+        and model_config.hf_config.hidden_size == 16384
+        and model_config.hf_config.intermediate_size == 53248
+        and model_config.hf_config.num_hidden_layers == 126
+        and model_config.hf_config.quantization_config["quant_method"] == "fbgemm_fp8"
+    ):
+        return True
+    return False
+
+
 class ModelRunner:
     def __init__(
         self,
@@ -119,6 +132,9 @@ class ModelRunner:
             seed=42,
             skip_tokenizer_init=True,
         )
+        if is_llama3_405b_fp8(self.model_config):
+            self.model_config.hf_config.num_key_value_heads = 8
+            vllm_model_config.hf_config.num_key_value_heads = 8
         self.dtype = vllm_model_config.dtype
         if self.model_config.model_overide_args is not None:
             vllm_model_config.hf_config.update(self.model_config.model_overide_args)
@@ -241,16 +257,20 @@ class ModelRunner:
             self.cuda_graph_runner = None
             return
 
-        logger.info(f"[gpu_id={self.gpu_id}] Capture cuda graph begin.")
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Capture cuda graph begin. This can take up to several minutes."
+        )
         batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 17)]
         self.cuda_graph_runner = CudaGraphRunner(
-            self, max_batch_size_to_capture=max(batch_size_list)
+            self,
+            max_batch_size_to_capture=max(batch_size_list),
+            use_torch_compile=self.server_args.enable_torch_compile,
         )
         try:
             self.cuda_graph_runner.capture(batch_size_list)
         except RuntimeError as e:
             raise Exception(
-                f"Capture cuda graph failed {e}. Possible solutions:\n"
+                f"Capture cuda graph failed: {e}. Possible solutions:\n"
                 f"1. disable cuda graph by --disable-cuda-graph\n"
                 f"2. set --mem-fraction-static to a smaller value\n"
                 f"Open an issue on GitHub with reproducible scripts if you need help.\n"
@@ -367,5 +387,39 @@ def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
     return model_arch_name_to_cls[model_arch]
 
 
+def get_original_weight(loaded_weight, head_dim):
+    n_kv_head = loaded_weight.shape[0] // (2 * head_dim)
+    dim = loaded_weight.shape[1]
+    for i in range(n_kv_head):
+        loaded_weight[i * head_dim : (i + 1) * head_dim, :] = loaded_weight[
+            2 * i * head_dim : (2 * i + 1) * head_dim, :
+        ]
+    original_kv_weight = loaded_weight[: n_kv_head * head_dim, :]
+    assert original_kv_weight.shape == (n_kv_head * head_dim, dim)
+    return original_kv_weight
+
+
+def get_weight_loader_srt(weight_loader):
+    def weight_loader_srt(
+        self,
+        param: Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: Optional[str] = None,
+    ):
+        if (
+            loaded_shard_id in ["k", "v"]
+            and loaded_weight.shape[0] == self.head_size * self.total_num_kv_heads * 2
+        ):
+            loaded_weight = get_original_weight(loaded_weight, self.head_size)
+
+        weight_loader(self, param, loaded_weight, loaded_shard_id)
+
+    return weight_loader_srt
+
+
 # Monkey patch model loader
 setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)
+original_weight_loader = QKVParallelLinear.weight_loader
+setattr(
+    QKVParallelLinear, "weight_loader", get_weight_loader_srt(original_weight_loader)
+)

sglang/srt/managers/controller/schedule_heuristic.py

@@ -14,7 +14,7 @@ class ScheduleHeuristic:
         tree_cache,
     ):
         if tree_cache.disable and schedule_heuristic == "lpm":
-            # LMP is not meaningless when tree cache is disabled.
+            # LMP is meaningless when the tree cache is disabled.
             schedule_heuristic = "fcfs"
 
         self.schedule_heuristic = schedule_heuristic
@@ -28,11 +28,16 @@ class ScheduleHeuristic:
             # longest prefix match
             forward_queue.sort(key=lambda x: -len(x.prefix_indices))
             return forward_queue
+        elif self.schedule_heuristic == "fcfs":
+            # first come first serve
+            return forward_queue
+        elif self.schedule_heuristic == "lof":
+            # longest output first
+            forward_queue.sort(key=lambda x: -x.sampling_params.max_new_tokens)
+            return forward_queue
         elif self.schedule_heuristic == "random":
             random.shuffle(forward_queue)
             return forward_queue
-        elif self.schedule_heuristic == "fcfs":
-            return forward_queue
         elif self.schedule_heuristic == "dfs-weight":
             last_node_to_reqs = defaultdict(list)
             for req in forward_queue:

sglang/srt/managers/controller/tp_worker.py

@@ -103,6 +103,9 @@ class ModelTpServer:
             if server_args.max_running_requests is None
             else server_args.max_running_requests
         )
+        self.max_running_requests = min(
+            self.max_running_requests, 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)
         )
@@ -113,13 +116,9 @@ class ModelTpServer:
             f"[gpu_id={self.gpu_id}] "
             f"max_total_num_tokens={self.max_total_num_tokens}, "
             f"max_prefill_tokens={self.max_prefill_tokens}, "
+            f"max_running_requests={self.max_running_requests}, "
             f"context_len={self.model_config.context_len}"
         )
-        if self.tp_rank == 0:
-            logger.info(
-                f"[gpu_id={self.gpu_id}] "
-                f"server_args: {server_args.print_mode_args()}"
-            )
 
         # Init cache
         self.tree_cache = RadixCache(
@@ -161,15 +160,12 @@ class ModelTpServer:
         assert (
             server_args.schedule_conservativeness >= 0
         ), "Invalid schedule_conservativeness"
-        self.new_token_ratio = min(
-            global_config.base_new_token_ratio * server_args.schedule_conservativeness,
-            1.0,
-        )
         self.min_new_token_ratio = min(
             global_config.base_min_new_token_ratio
             * server_args.schedule_conservativeness,
             1.0,
         )
+        self.new_token_ratio = self.min_new_token_ratio
         self.new_token_ratio_decay = global_config.new_token_ratio_decay
         self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
@@ -231,6 +227,7 @@ class ModelTpServer:
                         break
             else:
                 self.check_memory()
+                self.new_token_ratio = global_config.init_new_token_ratio
 
     def print_stats(self):
         num_used = self.max_total_num_tokens - (
@@ -539,9 +536,10 @@ class ModelTpServer:
         # Check if decode out of memory
         if not batch.check_decode_mem():
             old_ratio = self.new_token_ratio
-            self.new_token_ratio = min(old_ratio + self.new_token_ratio_recovery, 1.0)
 
-            retracted_reqs = batch.retract_decode()
+            retracted_reqs, new_token_ratio = batch.retract_decode()
+            self.new_token_ratio = new_token_ratio
+
             logger.info(
                 "decode out of memory happened, "
                 f"#retracted_reqs: {len(retracted_reqs)}, "

sglang/srt/memory_pool.py

@@ -11,6 +11,7 @@ class ReqToTokenPool:
     """A memory pool that maps a request to its token locations."""
 
     def __init__(self, size: int, max_context_len: int):
+        self.size = size
         self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
         self.req_to_token = torch.empty(
             (size, max_context_len), dtype=torch.int32, device="cuda"
@@ -57,9 +58,13 @@ class TokenToKVPool:
         # We also add one slot. This slot is used for writing dummy output from padded tokens.
         self.mem_state = torch.ones((self.size + 1,), dtype=torch.bool, device="cuda")
 
-        # [size, key/value, head_num, head_dim] for each layer
-        self.kv_data = [
-            torch.empty((size + 1, 2, head_num, head_dim), dtype=dtype, device="cuda")
+        # [size, head_num, head_dim] for each layer
+        self.k_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
+            for _ in range(layer_num)
+        ]
+        self.v_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
             for _ in range(layer_num)
         ]
 
@@ -71,10 +76,13 @@ class TokenToKVPool:
         self.clear()
 
     def get_key_buffer(self, layer_id: int):
-        return self.kv_data[layer_id][:, 0]
+        return self.k_buffer[layer_id]
 
     def get_value_buffer(self, layer_id: int):
-        return self.kv_data[layer_id][:, 1]
+        return self.v_buffer[layer_id]
+
+    def get_kv_buffer(self, layer_id: int):
+        return self.k_buffer[layer_id], self.v_buffer[layer_id]
 
     def available_size(self):
         return self.can_use_mem_size + len(self.prefetch_buffer)