diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 35a5487fef9e..63526e34a0e4 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 6e6d6fce436c..0063c522ffc1 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index f956705202ba..31c8c460a997 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index 670e72d37510..cfe81c0b0257 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/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,