From 0042fd977bd87f0305476b7acdbb367405c4d59a Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Fri, 17 Jul 2026 19:01:54 +0000 Subject: [PATCH 1/4] fix(trtllm-mla): avoid per-call multi-CTA KV counter buffer re-zeroing flashinfer's trtllm_batch_decode_with_kv_cache_mla instantiates a fresh runner per call, so its internal counter-buffer cache never survives across decode steps and it torch.zeros() a new one every call by default. Own a persistent buffer in TRTLLMMLABackend and DeepseekSparseAttnBackend and pass it via multi_ctas_kv_counter_buffer= instead. DSA's trtllm sparse-indexer path flattens every token in a prefill/ extend/verify chunk into its own decode-shaped batch entry, so its batch dimension can exceed max_running_requests; grow the buffer on demand there instead of assuming a fixed bound. --- .../srt/layers/attention/dsa_backend.py | 22 ++++++++++ .../layers/attention/trtllm_mla_backend.py | 41 +++++++++++++++++++ 2 files changed, 63 insertions(+) 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, From 121c901e65d729301ddea486f1f5b8e96de066f4 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Fri, 17 Jul 2026 23:11:44 +0000 Subject: [PATCH 2/4] fix --- .../kits/attention_unittest/attention_methods/dsa_attention.py | 1 + 1 file changed, 1 insertion(+) 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, From a84b52967a3f300bdda920e88846bade713088b3 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Sat, 18 Jul 2026 00:11:41 +0000 Subject: [PATCH 3/4] fix --- python/sglang/srt/model_executor/runner/flashinfer_autotune.py | 3 ++- .../kits/attention_unittest/attention_methods/mla_attention.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 46f4bba987f1..e5629ab7d5e7 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -174,9 +174,10 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool) maybe_skip_logits = autotune_dummy_run_mode() with torch.inference_mode(), autotune( # Autotuning mxfp8_gemm hits an IMA; skip it. + # Add trtllm MLA to skip autotune to fix a temporary bug. True, cache=str(autotune_cache), - skip_ops={"mxfp8_gemm"}, + skip_ops={"mxfp8_gemm", "trtllm_batch_decode_mla"}, ), maybe_skip_logits: yield torch.cuda.current_stream().wait_stream(mr.forward_stream) 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, From ed135a89cd3575ff7ea52f82b6d41d2a504293dd Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Sat, 18 Jul 2026 00:12:59 +0000 Subject: [PATCH 4/4] revert --- python/sglang/srt/model_executor/runner/flashinfer_autotune.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index e5629ab7d5e7..46f4bba987f1 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -174,10 +174,9 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool) maybe_skip_logits = autotune_dummy_run_mode() with torch.inference_mode(), autotune( # Autotuning mxfp8_gemm hits an IMA; skip it. - # Add trtllm MLA to skip autotune to fix a temporary bug. True, cache=str(autotune_cache), - skip_ops={"mxfp8_gemm", "trtllm_batch_decode_mla"}, + skip_ops={"mxfp8_gemm"}, ), maybe_skip_logits: yield torch.cuda.current_stream().wait_stream(mr.forward_stream)