sglang 0.1.18__py3-none-any.whl → 0.1.20__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -1,60 +1,42 @@
 """Radix attention."""
 
-import numpy as np
 import torch
+from flashinfer.cascade import merge_state
 from torch import nn
 
 from sglang.global_config import global_config
-from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
 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
 
 
 class RadixAttention(nn.Module):
     def __init__(
-        self, num_heads: int, head_dim: int, scaling: float, num_kv_heads: int,
-        layer_id: int, logit_cap: int = -1
+        self,
+        num_heads: int,
+        head_dim: int,
+        scaling: float,
+        num_kv_heads: int,
+        layer_id: int,
+        logit_cap: int = -1,
     ):
         super().__init__()
         self.tp_q_head_num = num_heads
         self.tp_k_head_num = num_kv_heads
         self.tp_v_head_num = num_kv_heads
         self.head_dim = head_dim
+        self.scaling = scaling
         self.layer_id = layer_id
 
-        assert np.allclose(scaling, 1.0 / (head_dim**0.5))
-
-        from sglang.srt.managers.controller.model_runner import global_server_args_dict
-
         if not global_server_args_dict.get("disable_flashinfer", False):
-            self.prefill_forward = self.prefill_forward_flashinfer
-            self.extend_forward = self.prefill_forward_flashinfer
+            self.extend_forward = self.extend_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer
-            # flashinfer now accepts float logit_cap argument
-            self.logit_cap = logit_cap if logit_cap > 0 else 0
         else:
-            self.prefill_forward = self.prefill_forward_triton
             self.extend_forward = self.extend_forward_triton
             self.decode_forward = self.decode_forward_triton
-            self.logit_cap = logit_cap
-
-    def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        o = torch.empty_like(q)
-
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.head_dim),
-            k,
-            v,
-            o.view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.start_loc,
-            input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-            self.logit_cap,
-        )
-        self.store_kv_cache(k, v, input_metadata)
 
-        return o
+        self.logit_cap = logit_cap if logit_cap is not None and logit_cap > 0 else 0
 
     def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -68,14 +50,15 @@ class RadixAttention(nn.Module):
             input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
             input_metadata.req_to_token_pool.req_to_token,
             input_metadata.req_pool_indices,
-            input_metadata.start_loc,
+            input_metadata.triton_start_loc,
             input_metadata.seq_lens,
-            input_metadata.prefix_lens,
+            input_metadata.triton_prefix_lens,
             input_metadata.extend_start_loc,
             input_metadata.extend_seq_lens,
-            input_metadata.max_seq_len,
-            input_metadata.max_extend_len,
-            self.logit_cap,
+            input_metadata.triton_max_seq_len,
+            input_metadata.triton_max_extend_len,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
@@ -91,39 +74,41 @@ class RadixAttention(nn.Module):
             o.view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.req_to_token_pool.req_to_token,
             input_metadata.req_pool_indices,
-            input_metadata.start_loc,
+            input_metadata.triton_start_loc,
             input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-            input_metadata.other_kv_index,
+            input_metadata.triton_max_seq_len,
             input_metadata.total_num_tokens,
-            self.logit_cap,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
 
-    def prefill_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        self.store_kv_cache(k, v, input_metadata)
-
+    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 input_metadata.no_prefix:
+        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.kv_data[self.layer_id],
                 causal=False,
+                sm_scale=self.scaling,
                 logits_soft_cap=self.logit_cap,
             )
 
-            from flashinfer.cascade import merge_state
             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()
 
@@ -135,6 +120,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],
+            sm_scale=self.scaling,
             logits_soft_cap=self.logit_cap,
         )
 
@@ -144,25 +130,13 @@ class RadixAttention(nn.Module):
         k = k.view(-1, self.tp_k_head_num, self.head_dim)
         v = v.view(-1, self.tp_v_head_num, self.head_dim)
 
-        if input_metadata.forward_mode == ForwardMode.PREFILL:
-            return self.prefill_forward(q, k, v, input_metadata)
-        elif input_metadata.forward_mode == ForwardMode.EXTEND:
+        if input_metadata.forward_mode == ForwardMode.EXTEND:
             return self.extend_forward(q, k, v, input_metadata)
         elif input_metadata.forward_mode == ForwardMode.DECODE:
             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)
-        if input_metadata.out_cache_loc is not None:
-            key_buffer[input_metadata.out_cache_loc] = cache_k
-            value_buffer[input_metadata.out_cache_loc] = cache_v
-        elif input_metadata.out_cache_cont_start is not None:
-            key_buffer[
-                input_metadata.out_cache_cont_start : input_metadata.out_cache_cont_end
-            ] = cache_k
-            value_buffer[
-                input_metadata.out_cache_cont_start : input_metadata.out_cache_cont_end
-            ] = cache_v
-        else:
-            raise RuntimeError()
+        value_buffer[input_metadata.out_cache_loc] = cache_v

sglang/srt/layers/token_attention.py

@@ -107,7 +107,6 @@ def _fwd_kernel_stage2(
     stride_obs,
     stride_oh,
     stride_req_to_token_b,
-    other_kv_index,  # To fix a NAN issue
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -138,7 +137,7 @@ def _fwd_kernel_stage2(
             + cur_batch_req_idx * stride_req_to_token_b
             + (start_n + offs_n),
             mask=(start_n + offs_n) < cur_batch_seq_len,
-            other=other_kv_index,
+            other=0,
         )
 
         qk = tl.load(
@@ -176,6 +175,7 @@ def _token_att_m_fwd(
     B_Start_Loc,
     B_Seqlen,
     max_len_in_batch,
+    sm_scale,
     logit_cap,
 ):
     BLOCK = 32
@@ -183,7 +183,6 @@ def _token_att_m_fwd(
     Lq, Lk = q.shape[-1], k_buffer.shape[-1]
     assert Lq == Lk
     assert Lk in {16, 32, 64, 128, 256}
-    sm_scale = 1.0 / (Lk**0.5)
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -250,7 +249,6 @@ def _token_softmax_reducev_fwd(
     b_req_idx,
     b_start_loc,
     b_seq_len,
-    other_kv_index,
 ):
     BLOCK = 64
     batch, head = b_seq_len.shape[0], logics.shape[0]
@@ -277,7 +275,6 @@ def _token_softmax_reducev_fwd(
             o.stride(0),
             o.stride(1),
             req_to_tokens.stride(0),
-            other_kv_index,
         )
         return
 
@@ -295,7 +292,6 @@ def _token_softmax_reducev_fwd(
         o.stride(0),
         o.stride(1),
         req_to_tokens.stride(0),
-        other_kv_index,
         kv_group_num=kv_group_num,
         BLOCK_DMODEL=v_buffer.shape[-1],
         BLOCK_N=BLOCK,
@@ -315,8 +311,8 @@ def token_attention_fwd(
     b_start_loc,
     b_seq_len,
     max_len_in_batch,
-    other_kv_index,
     total_num_tokens,
+    sm_scale,
     logit_cap=-1,
     att_m=None,
 ):
@@ -334,6 +330,7 @@ def token_attention_fwd(
         b_start_loc,
         b_seq_len,
         max_len_in_batch,
+        sm_scale,
         logit_cap,
     )
     _token_softmax_reducev_fwd(
@@ -344,5 +341,4 @@ def token_attention_fwd(
         b_req_idx,
         b_start_loc,
         b_seq_len,
-        other_kv_index,
     )

sglang/srt/managers/controller/cuda_graph_runner.py

@@ -0,0 +1,172 @@
+"""Run the model with cuda graph."""
+
+import bisect
+
+import torch
+from vllm.distributed.parallel_state import graph_capture
+
+from sglang.global_config import global_config
+from sglang.srt.layers.logits_processor import LogitProcessorOutput
+from sglang.srt.managers.controller.infer_batch import (
+    Batch, ForwardMode, InputMetadata, init_flashinfer_args
+)
+
+
+class CudaGraphRunner:
+    def __init__(self, model_runner, max_batch_size_to_capture):
+        self.model_runner = model_runner
+        self.graphs = {}
+        self.input_buffers = {}
+        self.output_buffers = {}
+        self.flashinfer_handlers = {}
+        self.graph_memory_pool = None
+
+        # Common inputs
+        self.max_bs = max_batch_size_to_capture
+        self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        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.max_bs,), dtype=torch.int32, device="cuda")
+        self.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+
+        # 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"
+        )
+        self.flashinfer_kv_indices = torch.zeros(
+            (self.max_bs * model_runner.model_config.context_len,), dtype=torch.int32, device="cuda"
+        )
+        self.flashinfer_kv_last_page_len = torch.ones(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
+
+    def can_run(self, batch_size):
+        return batch_size < self.max_bs
+
+    def capture(self, batch_size_list):
+        self.batch_size_list = batch_size_list
+        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
+
+        graph = torch.cuda.CUDAGraph()
+        stream = self.stream
+
+        # Common inputs
+        input_ids = self.input_ids[:bs]
+        req_pool_indices = self.req_pool_indices[:bs]
+        seq_lens = self.seq_lens[:bs]
+        position_ids_offsets = self.position_ids_offsets[:bs]
+        out_cache_loc = self.out_cache_loc[:bs]
+
+        # FlashInfer inputs
+        if not _grouped_size_compiled_for_decode_kernels(
+            self.model_runner.model_config.num_attention_heads // self.model_runner.tp_size,
+            self.model_runner.model_config.get_num_kv_heads(self.model_runner.tp_size),
+        ):
+            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],
+        )
+        init_flashinfer_args(
+            ForwardMode.DECODE,
+            self.model_runner,
+            req_pool_indices,
+            seq_lens,
+            None,
+            flashinfer_decode_wrapper,
+        )
+
+        # Run and capture
+        def run_once():
+            input_metadata = InputMetadata.create(
+                self.model_runner,
+                forward_mode=ForwardMode.DECODE,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens,
+                prefix_lens=None,
+                position_ids_offsets=position_ids_offsets,
+                out_cache_loc=out_cache_loc,
+                return_logprob=False,
+                top_logprobs_nums=0,
+                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
+            )
+
+        for _ in range(2):
+            run_once()
+
+        torch.cuda.synchronize()
+        with torch.cuda.graph(graph, pool=self.graph_memory_pool, stream=stream):
+            out = run_once()
+        torch.cuda.synchronize()
+        self.graph_memory_pool = graph.pool()
+        return graph, None, out, flashinfer_decode_wrapper
+
+    def replay(self, batch: Batch):
+        assert batch.out_cache_loc is not None
+        assert not batch.return_logprob
+        raw_bs = len(batch.reqs)
+
+        # Pad
+        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.out_cache_loc.zero_()
+
+        # Common inputs
+        self.input_ids[:raw_bs] = batch.input_ids
+        self.req_pool_indices[:raw_bs] = batch.req_pool_indices
+        self.seq_lens[:raw_bs] = batch.seq_lens
+        self.position_ids_offsets[:raw_bs] = batch.position_ids_offsets
+        self.out_cache_loc[:raw_bs] = batch.out_cache_loc
+
+        # FlashInfer inputs
+        init_flashinfer_args(
+            ForwardMode.DECODE,
+            self.model_runner,
+            self.req_pool_indices[:bs],
+            self.seq_lens[:bs],
+            None,
+            self.flashinfer_handlers[bs],
+        )
+
+        # Replay
+        self.graphs[bs].replay()
+        output = self.output_buffers[bs]
+
+        # Unpad
+        if bs == raw_bs:
+            return output
+        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,
+                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,
+            )
+        return output