sglang 0.1.21__py3-none-any.whl → 0.1.24__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -7,8 +7,8 @@ from torch import nn
 from sglang.global_config import global_config
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
-from sglang.srt.managers.controller.infer_batch import global_server_args_dict
 from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
+from sglang.srt.server import global_server_args_dict
 
 
 class RadixAttention(nn.Module):
@@ -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,7 +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):
-        key_buffer = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
-        key_buffer[input_metadata.out_cache_loc] = cache_k
-        value_buffer = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
-        value_buffer[input_metadata.out_cache_loc] = cache_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/layers/token_attention.py

@@ -5,8 +5,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.managers.controller.model_runner import global_server_args_dict
-from sglang.srt.utils import wrap_kernel_launcher
+from sglang.srt.server import global_server_args_dict
 
 if global_server_args_dict.get("attention_reduce_in_fp32", False):
     REDUCE_TRITON_TYPE = tl.float32
@@ -162,10 +161,6 @@ def _fwd_kernel_stage2(
     tl.store(out_ptrs, acc)
 
 
-cached_kernel_stage1 = None
-cached_kernel_stage2 = None
-
-
 def _token_att_m_fwd(
     q,
     k_buffer,
@@ -194,28 +189,6 @@ def _token_att_m_fwd(
     else:
         num_warps = 2
 
-    global cached_kernel_stage1
-    if cached_kernel_stage1:
-        cached_kernel_stage1(
-            grid,
-            num_warps,
-            q,
-            k_buffer,
-            sm_scale,
-            Req_to_tokens,
-            B_req_idx,
-            B_Start_Loc,
-            B_Seqlen,
-            att_out,
-            Req_to_tokens.stride(0),
-            q.stride(0),
-            q.stride(1),
-            k_buffer.stride(0),
-            k_buffer.stride(1),
-            att_out.stride(0),
-        )
-        return
-
     _fwd_kernel_stage1[grid](
         q,
         k_buffer,
@@ -238,7 +211,6 @@ def _token_att_m_fwd(
         num_warps=num_warps,
         num_stages=1,
     )
-    cached_kernel_stage1 = wrap_kernel_launcher(_fwd_kernel_stage1)
 
 
 def _token_softmax_reducev_fwd(
@@ -257,27 +229,6 @@ def _token_softmax_reducev_fwd(
 
     num_warps = 1
 
-    global cached_kernel_stage2
-    if cached_kernel_stage2:
-        cached_kernel_stage2(
-            grid,
-            num_warps,
-            logics,
-            v_buffer,
-            o,
-            req_to_tokens,
-            b_req_idx,
-            b_start_loc,
-            b_seq_len,
-            logics.stride(0),
-            v_buffer.stride(0),
-            v_buffer.stride(1),
-            o.stride(0),
-            o.stride(1),
-            req_to_tokens.stride(0),
-        )
-        return
-
     _fwd_kernel_stage2[grid](
         logics,
         v_buffer,
@@ -298,7 +249,6 @@ def _token_softmax_reducev_fwd(
         num_warps=num_warps,
         num_stages=3,
     )
-    cached_kernel_stage2 = wrap_kernel_launcher(_fwd_kernel_stage2)
 
 
 def token_attention_fwd(

sglang/srt/managers/controller/cuda_graph_runner.py

@@ -1,11 +1,14 @@
 """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.global_config import global_config
 from sglang.srt.layers.logits_processor import LogitProcessorOutput
 from sglang.srt.managers.controller.infer_batch import (
     Batch,
@@ -13,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 = {}
@@ -54,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
 
@@ -62,21 +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):
-        from flashinfer import BatchDecodeWithPagedKVCacheWrapper
-        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
-
+                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
 
@@ -129,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()
@@ -152,8 +192,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.zero_()
-            self.position_ids_offsets.fill_(1)
+            self.seq_lens.fill_(1)
+            self.position_ids_offsets.zero_()
             self.out_cache_loc.zero_()
 
         # Common inputs
@@ -183,14 +223,18 @@ class CudaGraphRunner:
         else:
             output = LogitProcessorOutput(
                 next_token_logits=output.next_token_logits[:raw_bs],
-                next_token_logprobs=output.next_token_logprobs[:raw_bs]
-                if output.next_token_logprobs is not None
-                else None,
+                next_token_logprobs=(
+                    output.next_token_logprobs[:raw_bs]
+                    if output.next_token_logprobs is not None
+                    else None
+                ),
                 normalized_prompt_logprobs=None,
                 prefill_token_logprobs=None,
                 prefill_top_logprobs=None,
-                decode_top_logprobs=output.decode_top_logprobs[:raw_bs]
-                if output.decode_top_logprobs is not None
-                else None,
+                decode_top_logprobs=(
+                    output.decode_top_logprobs[:raw_bs]
+                    if output.decode_top_logprobs is not None
+                    else None
+                ),
             )
         return output

sglang/srt/managers/controller/infer_batch.py

@@ -7,7 +7,9 @@ from typing import List, Union
 
 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
@@ -15,9 +17,6 @@ from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
-# Store some global server args
-global_server_args_dict = {}
-
 
 class ForwardMode(IntEnum):
     # Prefill a new sequence. This is deprecated now. "EXTEND" covers this case.
@@ -84,6 +83,15 @@ class Req:
         self.input_ids = None  # input_ids = origin_input_ids + output_ids
 
         # For incremental decoding
+        # ----- | --------- read_ids -------|
+        # ----- |   surr_ids  |
+        # xxxxx | xxxxxxxxxxx | xxxxxxxxxxx |
+        # ----- ^ ----------- ^ ----------- ^
+        # ----- 1 ----------- 2 ----------- 3
+        # 1: surr_offset
+        # 2: read_offset
+        # 3: last token
+        self.vid = 0  # version id to sync decode status with in detokenizer_manager
         self.decoded_text = ""
         self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
         self.read_offset = None
@@ -134,7 +142,7 @@ class Req:
         return self.finished_reason is not None
 
     # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
-    def init_detokenize_incrementally(self):
+    def init_incremental_detokenize(self):
         first_iter = self.surr_offset is None or self.read_offset is None
 
         if first_iter:
@@ -144,13 +152,11 @@ class Req:
             )
 
         all_ids = self.origin_input_ids_unpadded + self.output_ids
-        surr_ids = all_ids[self.surr_offset : self.read_offset]
-        read_ids = all_ids[self.surr_offset :]
-
-        return surr_ids, read_ids, len(all_ids)
+        return all_ids[self.surr_offset :], self.read_offset - self.surr_offset
 
-    def detokenize_incrementally(self, inplace: bool = True):
-        surr_ids, read_ids, num_all_tokens = self.init_detokenize_incrementally()
+    def get_next_inc_detokenization(self):
+        read_ids, read_offset = self.init_incremental_detokenize()
+        surr_ids = read_ids[:read_offset]
 
         surr_text = self.tokenizer.decode(
             surr_ids,
@@ -164,13 +170,7 @@ class Req:
         )
 
         if len(new_text) > len(surr_text) and not new_text.endswith("�"):
-            new_text = new_text[len(surr_text) :]
-            if inplace:
-                self.decoded_text += new_text
-                self.surr_offset = self.read_offset
-                self.read_offset = num_all_tokens
-
-            return True, new_text
+            return True, new_text[len(surr_text) :]
 
         return False, ""
 
@@ -272,6 +272,7 @@ class Batch:
     prefix_lens: torch.Tensor = None
     position_ids_offsets: torch.Tensor = None
     out_cache_loc: torch.Tensor = None
+    extend_num_tokens: int = None
 
     # For processing logprobs
     return_logprob: bool = False
@@ -282,10 +283,6 @@ class Batch:
     image_sizes: List[List[int]] = None
     image_offsets: List[int] = None
 
-    # Other arguments for control
-    output_ids: torch.Tensor = None
-    extend_num_tokens: int = None
-
     # Batched sampling params
     temperatures: torch.Tensor = None
     top_ps: torch.Tensor = None
@@ -327,6 +324,13 @@ class Batch:
         seq_lens = []
 
         req_pool_indices = self.req_to_token_pool.alloc(bs)
+
+        if req_pool_indices is None:
+            raise RuntimeError(
+                "Out of memory. "
+                "Please set a smaller number for `--max-running-requests`."
+            )
+
         req_pool_indices_cpu = req_pool_indices.cpu().numpy()
         for i in range(bs):
             flatten_input_ids.extend(input_ids[i])
@@ -398,10 +402,10 @@ class Batch:
         ).view(-1, 1)
         self.top_ps = torch.tensor(
             [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device
-        ).view(-1, 1)
+        )
         self.top_ks = torch.tensor(
             [r.sampling_params.top_k for r in reqs], dtype=torch.int, device=device
-        ).view(-1, 1)
+        )
         self.frequency_penalties = torch.tensor(
             [r.sampling_params.frequency_penalty for r in reqs],
             dtype=torch.float,
@@ -428,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),
@@ -440,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)
@@ -465,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 = []
@@ -499,7 +523,7 @@ class Batch:
                     cur_output_ids = req.output_ids
 
                     req.output_ids.extend(suffix_ids)
-                    decode_res, new_text = req.detokenize_incrementally(inplace=False)
+                    decode_res, new_text = req.get_next_inc_detokenization()
                     if not decode_res:
                         req.output_ids = cur_output_ids
                         continue
@@ -518,6 +542,9 @@ class Batch:
                         req.output_ids = cur_output_ids
                         continue
 
+                    # The decode status has diverged from detokenizer_manager
+                    req.vid += 1
+
                     # insert the old request into tree_cache
                     if req_pool_indices_cpu is None:
                         req_pool_indices_cpu = self.req_pool_indices.tolist()
@@ -659,20 +686,20 @@ class Batch:
 
         # TODO(lmzheng): apply penalty
         probs = torch.softmax(logits, dim=-1)
-        probs_sort, probs_idx = _top_p_top_k(probs, self.top_ps, self.top_ks)
-        try:
-            sampled_index = torch.multinomial(probs_sort, num_samples=1)
-        except RuntimeError as e:
-            warnings.warn(f"Ignore errors in sampling: {e}")
-            sampled_index = torch.ones(
-                probs_sort.shape[:-1] + (1,), dtype=torch.int64, device=probs.device
-            )
-        batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(
-            -1
+
+        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, success = top_k_top_p_sampling_from_probs(
+            probs, uniform_samples, self.top_ks, self.top_ps
         )
-        batch_next_token_probs = torch.gather(
-            probs_sort, dim=1, index=sampled_index
-        ).view(-1)
+
+        if torch.any(~success):
+            warnings.warn("Sampling failed, fallback to top_k=1 strategy")
+            probs = probs.masked_fill(torch.isnan(probs), 0.0)
+            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()
@@ -682,18 +709,7 @@ class Batch:
                         req.regex_fsm_state, batch_next_token_ids_cpu[i]
                     )
 
-        return batch_next_token_ids, batch_next_token_probs
-
-
-def _top_p_top_k(probs: torch.Tensor, top_ps: torch.Tensor, top_ks: torch.Tensor):
-    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
-    probs_sum = torch.cumsum(probs_sort, dim=-1)
-    probs_sort[(probs_sum - probs_sort) > top_ps] = 0.0
-    probs_sort[
-        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1) >= top_ks
-    ] = 0.0
-    probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
-    return probs_sort, probs_idx
+        return batch_next_token_ids
 
 
 @dataclass
@@ -731,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(
@@ -746,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,
@@ -754,6 +774,7 @@ class InputMetadata:
                 seq_lens,
                 prefix_lens,
                 model_runner.flashinfer_decode_wrapper,
+                use_ragged,
             )
 
         batch_size = len(req_pool_indices)
@@ -808,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:
@@ -828,16 +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)
@@ -870,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()
@@ -894,6 +920,7 @@ def init_flashinfer_args(
 
 
 def init_triton_args(forward_mode, seq_lens, prefix_lens):
+    """Init auxiliary variables for triton attention backend."""
     batch_size = len(seq_lens)
     max_seq_len = int(torch.max(seq_lens))
     start_loc = torch.zeros((batch_size,), dtype=torch.int32, device="cuda")