From ae23423b4693f67132256c5cf3bbb59f4249c1c1 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:44:26 -0700 Subject: [PATCH] Split TRTLLM MHA decode batches by KV sequence length (#34888) --- python/sglang/srt/environ.py | 4 + .../layers/attention/trtllm_mha_backend.py | 103 ++++++++++++++---- .../unittests/dense/test_trtllm_mha.py | 12 +- 3 files changed, 96 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2f5d560cd..4cc5e5015 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -916,6 +916,10 @@ class Envs: # None = standard attention. See https://arxiv.org/abs/2512.12087 SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR = EnvFloat(None) SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR = EnvFloat(None) + # Split TRTLLM-GEN decode attention into sorted, equal-size request groups. + # One preserves the default single-call path; values above one are useful + # for batches whose KV sequence lengths have a large spread. + SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS = EnvInt(1) # SM120 FlashMLA decode backend: "flashinfer" (default), "triton", or "torch". SGLANG_SM120_FLASHMLA_BACKEND = EnvStr("flashinfer") SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index d12d7266a..dfdebf508 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -271,6 +271,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): "trtllm_mha_default_kv_scale", lambda: torch.ones(1, dtype=torch.float32, device=self.device), ) + self.decode_seq_len_splits = envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.get() + if self.decode_seq_len_splits < 1: + raise ValueError( + "SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS must be at least 1, " + f"got {self.decode_seq_len_splits}" + ) def _check_decode_kv_access(self) -> None: supported_kinds = { @@ -1084,6 +1090,73 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): assert self.is_nvfp4_kvcache return self.kv_cache_quant_method.get_bmm_scales(layer.layer_id) + def _run_fixed_q_len_decode( + self, + query: torch.Tensor, + kv_cache, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + *, + bmm1_scale, + bmm2_scale, + window_left: int, + sinks: Optional[torch.Tensor], + q_len_per_req: int = 1, + kv_cache_sf=None, + ) -> torch.Tensor: + """Run decode, optionally sorting and splitting requests by KV length.""" + + def run_group(group_query, group_block_tables, group_seq_lens): + kwargs = {} + if q_len_per_req != 1: + kwargs["q_len_per_req"] = q_len_per_req + return flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query=group_query, + kv_cache=kv_cache, + workspace_buffer=self.workspace_buffer, + block_tables=group_block_tables, + seq_lens=group_seq_lens, + max_seq_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + window_left=window_left, + sinks=sinks, + skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), + out_dtype=self.q_data_type, + kv_cache_sf=kv_cache_sf, + multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, + **kwargs, + ) + + num_requests = seq_lens.shape[0] + num_splits = min(self.decode_seq_len_splits, num_requests) + if num_splits == 1: + return run_group(query, block_tables, seq_lens) + + order = torch.argsort(seq_lens) + query_by_request = query.view( + num_requests, q_len_per_req, query.shape[-2], query.shape[-1] + ) + output_by_request = torch.empty( + query_by_request.shape, + dtype=self.q_data_type, + device=query.device, + ) + for indices in torch.tensor_split(order, num_splits): + group_output = run_group( + query_by_request.index_select(0, indices).reshape( + -1, query.shape[-2], query.shape[-1] + ), + block_tables.index_select(0, indices), + seq_lens.index_select(0, indices), + ) + output_by_request.index_copy_( + 0, + indices, + group_output.view(-1, q_len_per_req, query.shape[-2], query.shape[-1]), + ) + return output_by_request.view(-1, query.shape[-2], query.shape[-1]) + def _get_nvfp4_decode_kv_cache(self, layer: RadixAttention) -> tuple[ tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor], @@ -1170,21 +1243,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): page_table = self._get_layer_page_table(layer, forward_batch) - o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( - query=q, - kv_cache=kv_cache, - workspace_buffer=self.workspace_buffer, - block_tables=page_table, - seq_lens=self.forward_metadata.cache_seqlens_int32, - max_seq_len=self.max_context_len, + o = self._run_fixed_q_len_decode( + q, + kv_cache, + page_table, + self.forward_metadata.cache_seqlens_int32, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, window_left=layer.sliding_window_size, sinks=attention_sink, - skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), - out_dtype=self.q_data_type, # model_runner.dtype kv_cache_sf=kv_cache_block_scales, - multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) if self.is_nvfp4_kvcache and o.dtype != self.q_data_type: o = o.to(self.q_data_type) @@ -1344,21 +1412,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) else: - o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( - query=q, - kv_cache=kv_cache, - workspace_buffer=self.workspace_buffer, - block_tables=page_table, - seq_lens=self.forward_metadata.cache_seqlens_int32, - max_seq_len=self.max_context_len, + o = self._run_fixed_q_len_decode( + q, + kv_cache, + page_table, + self.forward_metadata.cache_seqlens_int32, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, window_left=layer.sliding_window_size, sinks=attention_sink, - skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(), - out_dtype=self.q_data_type, q_len_per_req=self.forward_metadata.max_seq_len_q, - multi_ctas_kv_counter_buffer=self._multi_ctas_kv_counter_buffer, ) elif self.use_fmha_v2 and not cp_v2_active: # CP-v2 must go through cp_strategy.run_attention (per-shard diff --git a/test/registered/attention/unittests/dense/test_trtllm_mha.py b/test/registered/attention/unittests/dense/test_trtllm_mha.py index bbdfbcbb2..87b0fcc84 100644 --- a/test/registered/attention/unittests/dense/test_trtllm_mha.py +++ b/test/registered/attention/unittests/dense/test_trtllm_mha.py @@ -2,6 +2,7 @@ import unittest import torch +from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_flashinfer_available from sglang.srt.utils.common import ( @@ -176,8 +177,11 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): ) def test_projected_dense_decode_cases(self): - for case in self.DECODE_CASES: - with self.subTest(case=case.name, backend=case.backend): + for case_index, case in enumerate(self.DECODE_CASES): + splits = 2 if case_index == 0 else 1 + with self.subTest( + case=case.name, backend=case.backend + ), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(splits): run_dense_attention_case( self, case, @@ -226,7 +230,9 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): def test_runner_mode_frozen_kv_mtp_cuda_graph_runner_cases(self): for case in self.FROZEN_KV_MTP_RUNNER_CASES: - with self.subTest(case=case.name, backend=case.backend): + with self.subTest( + case=case.name, backend=case.backend + ), envs.SGLANG_TRTLLM_MHA_DECODE_SEQ_LEN_SPLITS.override(2): run_dense_frozen_kv_mtp_cuda_graph_runner_case( self, case,