Fix routed-experts device buffer overflow under DP attention (#26123)
This commit is contained in:
@@ -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__(
|
||||||
|
|||||||
Reference in New Issue
Block a user