From ab5a76995ae43c58d4c0103c3812a516ce4d258d Mon Sep 17 00:00:00 2001 From: Mohammad Angkad Date: Sat, 26 Sep 2026 05:19:42 +0000 Subject: [PATCH 1/2] Fix the pooled-indexer breakable prefill bridge under DP attention --- .../attention/dsa/kpool_prefill_cuda_graph.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) 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..88d8ee98c2a0 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") + if ( + result is None + or result.ndim != 2 + or 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 at most " + f"{(n, output.shape[1])}" + ) + num_rows = result.shape[0] # The following captured attention segment reads this stable padded buffer. - output[:n].copy_(result) - output[n:].fill_(-1) + output[:num_rows].copy_(result) + output[num_rows:].fill_(-1) def _kpool_indexer_prefill_capture_stub( From 6de44c6ecd374bc9ad14955e32f3555fc727e17e Mon Sep 17 00:00:00 2001 From: Mohammad Angkad Date: Sat, 26 Sep 2026 05:19:42 +0000 Subject: [PATCH 2/2] Bound the pooled-indexer bridge rows and cover DP + breakable prefill in CI --- .../layers/attention/dsa/kpool_prefill_cuda_graph.py | 12 ++++++------ test/registered/e2e/models/test_glm53_flash_b200.py | 2 ++ 2 files changed, 8 insertions(+), 6 deletions(-) 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 88d8ee98c2a0..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,21 +37,21 @@ def _kpool_indexer_prefill_with_output( ) if not return_indices: return + num_logical = sum(forward_batch.extend_seq_lens_cpu) if ( result is None or result.ndim != 2 - or result.shape[0] > n + 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 at most " - f"{(n, output.shape[1])}" + f"{None if result is None else tuple(result.shape)}, expected " + f"between {num_logical} and {n} rows of width {output.shape[1]}" ) - num_rows = result.shape[0] # The following captured attention segment reads this stable padded buffer. - output[:num_rows].copy_(result) - output[num_rows:].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", ]