diff --git a/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py
index 8fb92d09f1b5..385607af6fe2 100644
--- a/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py
+++ b/python/sglang/srt/layers/attention/dsa/kpool_prefill_cuda_graph.py
@@ -37,11 +37,21 @@ def _kpool_indexer_prefill_with_output(
     )
     if not return_indices:
         return
-    if result is None or result.shape != (n, output.shape[1]):
-        raise ValueError("Pooled-indexer prefill returned an unexpected top-k shape")
+    num_logical = sum(forward_batch.extend_seq_lens_cpu)
+    if (
+        result is None
+        or result.ndim != 2
+        or not num_logical <= result.shape[0] <= n
+        or result.shape[1] != output.shape[1]
+    ):
+        raise ValueError(
+            "Pooled-indexer prefill returned an unexpected top-k shape: got "
+            f"{None if result is None else tuple(result.shape)}, expected "
+            f"between {num_logical} and {n} rows of width {output.shape[1]}"
+        )
     # The following captured attention segment reads this stable padded buffer.
-    output[:n].copy_(result)
-    output[n:].fill_(-1)
+    output[:num_logical].copy_(result[:num_logical])
+    output[num_logical:].fill_(-1)
 
 
 def _kpool_indexer_prefill_capture_stub(
diff --git a/test/registered/e2e/models/test_glm53_flash_b200.py b/test/registered/e2e/models/test_glm53_flash_b200.py
index d11920d884fd..38d2a1ae1e97 100644
--- a/test/registered/e2e/models/test_glm53_flash_b200.py
+++ b/test/registered/e2e/models/test_glm53_flash_b200.py
@@ -108,6 +108,8 @@ class TestGLM53FlashB200HighThroughput(
         "--enable-dp-attention",
         "--dp-size",
         "4",
+        "--cuda-graph-backend-prefill",
+        "breakable",
     ]