Avoid TRTLLM prefill output copy (#33306)

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