sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/layers/attention/flashmla_backend.py

@@ -19,7 +19,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
-    from sglang.srt.speculative.spec_info import SpecInfo
+    from sglang.srt.speculative.spec_info import SpecInput
 
 
 # FlashMLA only supports pagesize=64
@@ -187,7 +187,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[SpecInput],
     ):
         if forward_mode.is_decode_or_idle():
             max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
@@ -201,9 +201,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
                 self.req_to_token.stride(0),
                 self.cuda_graph_kv_indices.stride(0),
             )
+            num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
             mla_metadata, num_splits = get_mla_metadata(
                 seq_lens.to(torch.int32),
-                self.num_q_heads,
+                num_q_heads,
                 1,
             )
             self.cuda_graph_mla_metadata.copy_(mla_metadata)
@@ -257,7 +258,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
     ):
 
@@ -275,9 +276,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
                 self.req_to_token.stride(0),
                 self.cuda_graph_kv_indices.stride(0),
             )
+            num_q_heads = self.num_q_heads * (self.num_draft_tokens or 1)
             mla_metadata, num_splits = get_mla_metadata(
                 seq_lens.to(torch.int32),
-                self.num_q_heads,
+                num_q_heads,
                 1,
             )
             self.cuda_graph_mla_metadata.copy_(mla_metadata)

sglang/srt/layers/attention/hybrid_attn_backend.py

@@ -3,10 +3,11 @@ from typing import Optional, Union
 import torch
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.nsa.nsa_indexer import BaseIndexerMetadata
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
-from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
+from sglang.srt.speculative.spec_info import SpecInput
 
 
 class HybridAttnBackend(AttentionBackend):
@@ -21,18 +22,46 @@ class HybridAttnBackend(AttentionBackend):
         self.model_runner = model_runner
         self.prefill_backend = prefill_backend
         self.decode_backend = decode_backend
+        self.data_type = model_runner.kv_cache_dtype
 
-    def init_forward_metadata(self, forward_batch: ForwardBatch):
-        if forward_batch.forward_mode.is_decode_or_idle():
-            self.decode_backend.init_forward_metadata(forward_batch)
+    def _select_backend(self, forward_mode: ForwardMode) -> AttentionBackend:
+        """
+        Select the appropriate attention backend based on the forward mode.
+
+        Args:
+            forward_mode: The current forward mode indicating the operation type
+
+        Returns:
+            The selected attention backend (prefill or decode)
+
+        Note:
+            - decode_or_idle: Always uses decode backend
+            - target_verify or draft_extend: Uses decode backend if speculative_attention_mode is "decode", otherwise prefill backend
+            - prefill: Always uses prefill backend
+        """
+        if forward_mode.is_decode_or_idle():
+            return self.decode_backend
+        elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
+            return (
+                self.decode_backend
+                if self.model_runner.server_args.speculative_attention_mode == "decode"
+                else self.prefill_backend
+            )
         else:
-            self.prefill_backend.init_forward_metadata(forward_batch)
+            return self.prefill_backend
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        backend = self._select_backend(forward_batch.forward_mode)
+        backend.init_forward_metadata(forward_batch)
 
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
         self.decode_backend.init_cuda_graph_state(max_bs, max_num_tokens)
-        if self.model_runner.server_args.speculative_algorithm is not None:
-            # When speculative decoding is enabled, we also need to initialize the
-            # prefill backend's cuda graph state to support target_verify.
+        if (
+            self.model_runner.server_args.speculative_algorithm is not None
+            and self.model_runner.server_args.speculative_attention_mode == "prefill"
+        ):
+            # When speculative decoding is enabled, we need to initialize the backend
+            # that will be used for target_verify.
             self.prefill_backend.init_cuda_graph_state(max_bs, max_num_tokens)
 
     def init_forward_metadata_capture_cuda_graph(
@@ -43,28 +72,18 @@ class HybridAttnBackend(AttentionBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
     ):
-        if forward_mode.is_decode_or_idle():
-            self.decode_backend.init_forward_metadata_capture_cuda_graph(
-                bs,
-                num_tokens,
-                req_pool_indices,
-                seq_lens,
-                encoder_lens,
-                forward_mode,
-                spec_info,
-            )
-        else:
-            self.prefill_backend.init_forward_metadata_capture_cuda_graph(
-                bs,
-                num_tokens,
-                req_pool_indices,
-                seq_lens,
-                encoder_lens,
-                forward_mode,
-                spec_info,
-            )
+        backend = self._select_backend(forward_mode)
+        backend.init_forward_metadata_capture_cuda_graph(
+            bs,
+            num_tokens,
+            req_pool_indices,
+            seq_lens,
+            encoder_lens,
+            forward_mode,
+            spec_info,
+        )
 
     def init_forward_metadata_replay_cuda_graph(
         self,
@@ -74,31 +93,20 @@ class HybridAttnBackend(AttentionBackend):
         seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        spec_info: Optional[SpecInput],
         seq_lens_cpu: Optional[torch.Tensor],
     ):
-        if forward_mode.is_decode_or_idle():
-            self.decode_backend.init_forward_metadata_replay_cuda_graph(
-                bs,
-                req_pool_indices,
-                seq_lens,
-                seq_lens_sum,
-                encoder_lens,
-                forward_mode,
-                spec_info,
-                seq_lens_cpu,
-            )
-        else:
-            self.prefill_backend.init_forward_metadata_replay_cuda_graph(
-                bs,
-                req_pool_indices,
-                seq_lens,
-                seq_lens_sum,
-                encoder_lens,
-                forward_mode,
-                spec_info,
-                seq_lens_cpu,
-            )
+        backend = self._select_backend(forward_mode)
+        backend.init_forward_metadata_replay_cuda_graph(
+            bs,
+            req_pool_indices,
+            seq_lens,
+            seq_lens_sum,
+            encoder_lens,
+            forward_mode,
+            spec_info,
+            seq_lens_cpu,
+        )
 
     def get_cuda_graph_seq_len_fill_value(self):
         return self.decode_backend.get_cuda_graph_seq_len_fill_value()
@@ -127,6 +135,13 @@ class HybridAttnBackend(AttentionBackend):
         save_kv_cache: bool = True,
         **kwargs,
     ):
-        return self.prefill_backend.forward_extend(
+        backend = self._select_backend(forward_batch.forward_mode)
+        return backend.forward_extend(
             q, k, v, layer, forward_batch, save_kv_cache, **kwargs
         )
+
+    def get_indexer_metadata(
+        self, layer_id: int, forward_batch: ForwardBatch
+    ) -> Optional[BaseIndexerMetadata]:
+        backend = self._select_backend(forward_batch.forward_mode)
+        return backend.get_indexer_metadata(layer_id, forward_batch)