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

This commit is contained in:
Bruce Changlong Xu
2026-05-30 19:11:06 -07:00
committed by GitHub
parent f220c72929
commit 376635c1e3
@@ -59,12 +59,14 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
num_layers = model_config.hf_text_config.num_hidden_layers num_layers = model_config.hf_text_config.num_hidden_layers
server_args = get_global_server_args() server_args = get_global_server_args()
# FIXME: spec decoding is not accounted for here. The device buffer can # Scale by dp_size so the buffer covers the full DP-concatenated batch.
# overflow when max_running_requests * num_verify_tokens exceeds # _get_local_slice indexes into [attention_dp_rank * cuda_graph_batch, ...)
# chunked_prefill_size * dp_size. # 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( max_batch_size = max(
server_args.chunked_prefill_size * server_args.dp_size, server_args.chunked_prefill_size * server_args.dp_size,
max_running_requests, max_running_requests * server_args.dp_size,
) )
super().__init__( super().__init__(