Fix routed-experts device buffer overflow under DP attention (#26123)

2026-05-30 19:11:06 -07:00
parent f220c72929
commit 376635c1e3
1 changed files with 6 additions and 4 deletions
@@ -59,12 +59,14 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
        num_layers = model_config.hf_text_config.num_hidden_layers

        server_args = get_global_server_args()
-        # FIXME: spec decoding is not accounted for here. The device buffer can
-        # overflow when max_running_requests * num_verify_tokens exceeds
-        # chunked_prefill_size * dp_size.
+        # Scale by dp_size so the buffer covers the full DP-concatenated batch.
+        # _get_local_slice indexes into [attention_dp_rank * cuda_graph_batch, ...)
+        # and otherwise overflows on dp_rank > 0 when max_running_requests >
+        # chunked_prefill_size.
+        # FIXME: spec decoding's num_verify_tokens is still not accounted for.
        max_batch_size = max(
            server_args.chunked_prefill_size * server_args.dp_size,
-            max_running_requests,
+            max_running_requests * server_args.dp_size,
        )

        super().__init__(