sglang 0.4.0__py3-none-any.whl → 0.4.0.post2__py3-none-any.whl

sglang/srt/layers/attention/triton_backend.py

@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.srt.layers.attention import AttentionBackend
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 if TYPE_CHECKING:
@@ -35,10 +34,8 @@ class TritonAttnBackend(AttentionBackend):
                 model_runner.model_config.num_attention_heads // model_runner.tp_size
             )
 
-        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
-            self.reduce_dtype = torch.float32
-        else:
-            self.reduce_dtype = torch.float16
+        self.num_kv_splits = model_runner.server_args.triton_attention_num_kv_splits
+        self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
 
         self.forward_metadata = None
 
@@ -50,23 +47,23 @@ class TritonAttnBackend(AttentionBackend):
         """Init auxiliary variables for triton attention backend."""
 
         if forward_batch.forward_mode.is_decode():
-            start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
-            start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)
-
-            total_num_tokens = forward_batch.seq_lens_sum
             attn_logits = torch.empty(
-                (self.num_head, total_num_tokens),
-                dtype=self.reduce_dtype,
+                (
+                    forward_batch.batch_size,
+                    self.num_head,
+                    self.num_kv_splits,
+                    self.v_head_dim + 1,
+                ),
+                dtype=torch.float32,
                 device=self.device,
             )
 
-            max_seq_len = torch.max(forward_batch.seq_lens).item()
             max_extend_len = None
         else:
-            start_loc = attn_logits = max_seq_len = None
+            attn_logits = None
             max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
 
-        self.forward_metadata = start_loc, attn_logits, max_seq_len, max_extend_len
+        self.forward_metadata = attn_logits, max_extend_len
 
     def init_cuda_graph_state(self, max_bs: int):
         self.cuda_graph_max_total_num_tokens = max_bs * self.cuda_graph_max_seq_len
@@ -75,11 +72,8 @@ class TritonAttnBackend(AttentionBackend):
             (max_bs,), dtype=torch.int32, device=self.device
         )
         self.cuda_graph_attn_logits = torch.empty(
-            (
-                self.num_head,
-                self.cuda_graph_max_total_num_tokens,
-            ),
-            dtype=self.reduce_dtype,
+            (max_bs, self.num_head, self.num_kv_splits, self.v_head_dim + 1),
+            dtype=torch.float32,
             device="cuda",
         )
 
@@ -92,9 +86,7 @@ class TritonAttnBackend(AttentionBackend):
     ):
         # NOTE: encoder_lens expected to be zeros or None
         self.forward_metadata = (
-            self.cuda_graph_start_loc,
             self.cuda_graph_attn_logits,
-            self.cuda_graph_max_seq_len,
             None,
         )
 
@@ -114,7 +106,13 @@ class TritonAttnBackend(AttentionBackend):
         return 1
 
     def forward_extend(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         # TODO: reuse the buffer across layers
         if layer.qk_head_dim != layer.v_head_dim:
@@ -122,11 +120,12 @@ class TritonAttnBackend(AttentionBackend):
         else:
             o = torch.empty_like(q)
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
 
-        start_loc, attn_logits, max_seq_len, max_extend_len = self.forward_metadata
+        _, max_extend_len = self.forward_metadata
         self.extend_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
             k.contiguous(),
@@ -146,7 +145,13 @@ class TritonAttnBackend(AttentionBackend):
         return o
 
     def forward_decode(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         # During torch.compile, there is a bug in rotary_emb that causes the
         # output value to have a 3D tensor shape. This reshapes the output correctly.
@@ -158,11 +163,12 @@ class TritonAttnBackend(AttentionBackend):
         else:
             o = torch.empty_like(q)
 
-        start_loc, attn_logits, max_seq_len, max_extend_len = self.forward_metadata
+        attn_logits, _ = self.forward_metadata
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
 
         self.decode_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
@@ -171,10 +177,9 @@ class TritonAttnBackend(AttentionBackend):
             o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
             forward_batch.req_to_token_pool.req_to_token,
             forward_batch.req_pool_indices,
-            start_loc,
             forward_batch.seq_lens,
             attn_logits,
-            max_seq_len,
+            self.num_kv_splits,
             layer.scaling,
             layer.logit_cap,
         )