[attn backend] avoid initing parent class's workspace buffer (#25321)

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -197,6 +197,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         skip_prefill: bool = False,
         kv_indptr_buf: Optional[torch.Tensor] = None,
         q_indptr_decode_buf: Optional[torch.Tensor] = None,
+        skip_init_workspace_buffer: bool = False,
     ):
         super().__init__()
 
@@ -204,6 +205,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         self.max_context_len = model_runner.model_config.context_len
         self.device = model_runner.device
         self.skip_prefill = skip_prefill
+        self.skip_init_workspace_buffer = skip_init_workspace_buffer
         self.enable_chunk_kv = (
             not skip_prefill
             and get_global_server_args().disaggregation_mode != "decode"
@@ -213,6 +215,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         self.page_size = model_runner.page_size
 
         # Allocate buffers
+        if skip_init_workspace_buffer:
+            self.workspace_buffer = None
+        else:
             global global_workspace_buffer
             if global_workspace_buffer is None:
                 # different from flashinfer zero_init_global_workspace_buffer
@@ -243,6 +248,15 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         else:
             self.q_indptr_decode = q_indptr_decode_buf
 
+        if skip_init_workspace_buffer:
+            self.fmha_backend = None
+            self.prefill_wrapper_ragged = None
+            self.prefill_wrapper_paged = None
+            self.prefill_wrapper_verify = None
+            self.decode_wrapper = None
+            self.indices_updater_prefill = None
+            self.indices_updater_decode = None
+        else:
             if is_sm100_supported():
                 self.fmha_backend = "cutlass"
             else:
@@ -274,7 +288,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
                     model_runner, self
                 )
                 if self.enable_chunk_kv:
-                self.mha_chunk_kv_cache = FlashInferMhaChunkKVRunner(model_runner, self)
+                    self.mha_chunk_kv_cache = FlashInferMhaChunkKVRunner(
+                        model_runner, self
+                    )
 
             self.indices_updater_decode = FlashInferMLAIndicesUpdaterDecode(
                 model_runner, self

python/sglang/srt/layers/attention/tokenspeed_mla_backend.py

@@ -92,6 +92,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
             skip_prefill,
             kv_indptr_buf,
             q_indptr_decode_buf,
+            skip_init_workspace_buffer=True,
         )
 
         if self.data_type != torch.float8_e4m3fn:

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -263,12 +263,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         skip_prefill: bool = False,
         kv_indptr_buf: Optional[torch.Tensor] = None,
         q_indptr_decode_buf: Optional[torch.Tensor] = None,
+        skip_init_workspace_buffer: bool = False,
     ):
         super().__init__(
             model_runner,
             skip_prefill,
             kv_indptr_buf,
             q_indptr_decode_buf,
+            skip_init_workspace_buffer=True,
         )
 
         config = model_runner.model_config
@@ -294,6 +296,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
 
         # Workspace allocation
         self.workspace_size = DEFAULT_WORKSPACE_SIZE_MB * 1024 * 1024
+        if skip_init_workspace_buffer:
+            self.workspace_buffer = None
+        else:
             global global_zero_init_workspace_buffer
             if global_zero_init_workspace_buffer is None:
                 global_zero_init_workspace_buffer = torch.zeros(
@@ -378,6 +383,10 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
 
         return block_kv_indices
 
+    def init_mha_chunk_metadata(self, forward_batch: "ForwardBatch") -> None:
+        """Skip parent's flashinfer wrapper plan()."""
+        return None
+
     def init_cuda_graph_state(
         self,
         max_bs: int,
@@ -672,9 +681,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
         else:
             return super().init_forward_metadata(forward_batch)
 
-    def init_mha_chunk_metadata(self, forward_batch: ForwardBatch):
-        super().init_mha_chunk_metadata(forward_batch, disable_flashinfer_ragged=True)
-
     def pad_draft_extend_query(
         self,
         q: torch.Tensor,