diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index b09f22664..20881c205 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -1302,6 +1302,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): *, cu_seqlens_kv, use_zigzag_page_table=False, + out=None, ): block_tables = page_table if use_zigzag_page_table: @@ -1327,6 +1328,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): window_left=layer.sliding_window_size, sinks=attention_sink, skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_PREFILL_THRESHOLD_SCALE_FACTOR.get(), + out=out, out_dtype=self.q_data_type, ) @@ -1341,12 +1343,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): attention_backend=CPAttentionBackendKind.TRTLLM_MHA, ) else: + out = forward_batch._attn_output + if out is not None: + out = out.view_as(q) o = _trtllm_context_attn( q, self.forward_metadata.cu_seqlens_q, self.forward_metadata.cache_seqlens_int32, self.forward_metadata.max_seq_len_q, cu_seqlens_kv=self.forward_metadata.cu_seqlens_k, + out=out, ) return o.view(-1, layer.tp_q_head_num * layer.head_dim) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index fc1cdb8c9..a83c090b7 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -553,6 +553,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): num_token_non_padded_cpu: int = None # === Runtime-filled (set during the forward pass / cuda graph / managers; not at construction) === + # Preallocated piecewise-graph attention output, set by RadixAttention. + _attn_output: Optional[torch.Tensor] = None + # For logits and logprobs post processing next_token_logits_buffer: torch.Tensor = None temperature: torch.Tensor = None diff --git a/test/registered/attention/unittests/dense/test_trtllm_mha.py b/test/registered/attention/unittests/dense/test_trtllm_mha.py index ac65a048a..36c0bd4a1 100644 --- a/test/registered/attention/unittests/dense/test_trtllm_mha.py +++ b/test/registered/attention/unittests/dense/test_trtllm_mha.py @@ -28,6 +28,9 @@ from sglang.test.kits.attention_unittest.runner_modes.speculative_draft_runner i run_dense_eagle_draft_cuda_graph_runner_case, run_dense_frozen_kv_mtp_cuda_graph_runner_case, ) +from sglang.test.kits.attention_unittest.runner_modes.split_op_runner import ( + run_dense_split_op_extend_case, +) register_cuda_ci(est_time=20, stage="base-b", runner_config="4-gpu-b200") register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-large") @@ -125,6 +128,21 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): prefix_lens=(31, 32), ), ) + SPLIT_OP_CASES = ( + ( + DenseAttentionCase( + name="runner_split_op_trtllm_mha_extend_ragged", + backend="trtllm_mha", + forward_mode=ForwardMode.EXTEND, + num_heads=4, + num_kv_heads=4, + page_size=16, + prefix_lens=(0, 8, 16), + extend_lens=(15, 8, 1), + ), + 32, + ), + ) # EAGLE draft CG runner — chain only (topk=1). trtllm_mha is constrained # to topk=1 via `trtllm_mha_backend.py:459,492` so tree-mode tests don't @@ -182,6 +200,23 @@ class TestTRTLLMMHADenseAttentionBackendCorrectness(CustomTestCase): hidden_size=self.HIDDEN_SIZE, ) + @unittest.skipUnless( + is_sm100_supported(), "TRT-LLM context attention requires SM100" + ) + def test_runner_mode_split_op_extend_cases(self): + for case, static_num_tokens in self.SPLIT_OP_CASES: + for breakable in (False, True): + runner = "bcg" if breakable else "pcg" + with self.subTest(case=case.name, backend=case.backend, runner=runner): + run_dense_split_op_extend_case( + self, + case, + breakable=breakable, + static_num_tokens=static_num_tokens, + head_dim=self.HEAD_DIM, + hidden_size=self.HIDDEN_SIZE, + ) + def test_runner_mode_eagle_draft_cuda_graph_runner_cases(self): for case, topk, num_draft_tokens in self.EAGLE_DRAFT_RUNNER_CASES: with self.subTest(case=case.name, backend=case.backend, topk=topk):