Fix the perf regression of zeroing of TensorRT-LLM MLA (and sparse) workspace

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

@@ -52,6 +52,10 @@
     pad_dsa_cache_seqlens,
     should_use_dsa_fused_topk,
 )
+from sglang.srt.layers.attention.trtllm_mla_backend import (
+    grow_multi_ctas_kv_counter_buffer_if_needed,
+    make_persistent_multi_ctas_kv_counter_buffer,
+)
 from sglang.srt.layers.utils.cp_utils import (
     cp_all_gather_rerange_output,
     cp_split_and_rebuild_position,
@@ -458,8 +462,16 @@ def __init__(
                     device=model_runner.device,
                 ),
             )
+            self._multi_ctas_kv_counter_buffer = (
+                make_persistent_multi_ctas_kv_counter_buffer(
+                    torch.device(self.device),
+                    self.num_q_heads,
+                    max_batch_size=model_runner.max_running_requests,
+                )
+            )
         else:
             self.workspace_buffer = None
+            self._multi_ctas_kv_counter_buffer = None
 
     def _make_aiter_dsa_decode_metadata_buffer(
         self,
@@ -2719,6 +2731,15 @@ def _forward_trtllm(
         batch_size = page_table_1.shape[0]
         _, num_heads, head_dim = q_all.shape
 
+        self._multi_ctas_kv_counter_buffer = (
+            grow_multi_ctas_kv_counter_buffer_if_needed(
+                self._multi_ctas_kv_counter_buffer,
+                torch.device(self.device),
+                self.num_q_heads,
+                batch_size,
+            )
+        )
+
         q = q_all.view(batch_size, 1, num_heads, head_dim)
         kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim)
         block_tables = page_table_1.unsqueeze(1)
@@ -2747,6 +2768,7 @@ def _forward_trtllm(
             bmm1_scale=bmm1_scale,
             backend="trtllm-gen",
             skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
+            multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer,
         )
 
         return out

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

@@ -61,6 +61,35 @@
 # compute the LCM with other padding constraints.
 TRTLLM_BLOCK_CONSTRAINT = 128
 
+TRTLLM_MLA_MAX_BATCH_SIZE = 8192
+
+
+def _multi_ctas_kv_counter_bytes(
+    device: torch.device, num_q_heads: int, batch_size: int
+) -> int:
+    sm_count = flashinfer.utils.get_device_sm_count(device)
+    return flashinfer.utils.get_trtllm_gen_multi_ctas_kv_counter_bytes(
+        batch_size, num_q_heads, sm_count
+    )
+
+
+def make_persistent_multi_ctas_kv_counter_buffer(
+    device: torch.device, num_q_heads: int, max_batch_size: int
+) -> torch.Tensor:
+    num_bytes = _multi_ctas_kv_counter_bytes(
+        device, num_q_heads, max(TRTLLM_MLA_MAX_BATCH_SIZE, max_batch_size)
+    )
+    return torch.zeros(num_bytes, dtype=torch.uint8, device=device)
+
+
+def grow_multi_ctas_kv_counter_buffer_if_needed(
+    buffer: torch.Tensor, device: torch.device, num_q_heads: int, batch_size: int
+) -> torch.Tensor:
+    required_bytes = _multi_ctas_kv_counter_bytes(device, num_q_heads, batch_size)
+    if buffer.numel() >= required_bytes:
+        return buffer
+    return torch.zeros(required_bytes, dtype=torch.uint8, device=device)
+
 
 def _quantize_fp8_qkv(q, k, v, layer):
     q = q.to(torch.float8_e4m3fn)
@@ -189,6 +218,14 @@ def __init__(
                 ),
             )
 
+        self._multi_ctas_kv_counter_buffer = (
+            make_persistent_multi_ctas_kv_counter_buffer(
+                torch.device(self.device),
+                self.num_q_heads,
+                max_batch_size=model_runner.max_running_requests,
+            )
+        )
+
         # CUDA graph state
         self.decode_cuda_graph_metadata = {}
         self.decode_cuda_graph_kv_indices = None
@@ -637,6 +674,10 @@ def _run_decode_kernel(
             seq_lens if seq_lens.dtype == torch.int32 else seq_lens.to(torch.int32)
         )
         extra_kwargs = {"backend": self.backend} if self.backend != "trtllm-gen" else {}
+        if self.backend == "trtllm-gen":
+            extra_kwargs["multi_ctas_kv_counter_buffer"] = (
+                self._multi_ctas_kv_counter_buffer
+            )
         return flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
             query=query,
             kv_cache=kv_cache,

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -354,6 +354,7 @@ def __init__(
             triton_attention_split_tile_size=None,
         )
         self.server_args = self._server_args_override.install()
+        self.max_running_requests = pool_batch_size
         self.req_to_token_pool = ReqToTokenPool(
             size=pool_batch_size,
             max_context_len=max_context_len,

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -282,6 +282,7 @@ def __init__(
             triton_attention_split_tile_size=None,
         )
         self.server_args = self._server_args_override.install()
+        self.max_running_requests = pool_batch_size
         self.req_to_token_pool = ReqToTokenPool(
             size=pool_batch_size,
             max_context_len=max_context_len,