Avoid TRTLLM prefill output copy (#33306)

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)

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

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):