Revert "[attn backend] avoid initing parent class's workspace buffer" (#25488)

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

@@ -197,7 +197,6 @@ 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__()
 
@@ -205,7 +204,6 @@ 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"
@@ -215,9 +213,6 @@ 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
@@ -248,15 +243,6 @@ 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:
@@ -288,9 +274,7 @@ 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,7 +92,6 @@ 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,14 +263,12 @@ 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
@@ -296,9 +294,6 @@ 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(
@@ -383,10 +378,6 @@ 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,
@@ -681,6 +672,9 @@ 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,