Repository: vllm-project/vllm
File: vllm/v1/worker/gpu/dp_utils.py (Python)

Code before:


def get_cudagraph_and_dp_padding(
    num_tokens: int,
    cudagraph_size: int | None,
    cudagraph_runtime_mode: int,
    dp_size: int,
    dp_rank: int,
) -> tuple[int, torch.Tensor | None, int]:
    if dp_size == 1:
        if cudagraph_size is not None:
            return cudagraph_size, None, cudagraph_runtime_mode
        else:
            return num_tokens, None, cudagraph_runtime_mode

    # Convert None to -1 for sync (indicates no cudagraph available)
    if num_tokens == 0:
        cudagraph_size = 0
    elif cudagraph_size is None:
        cudagraph_size = -1

    (
        num_tokens_across_dp,
        cudagraph_size_across_dp,
        cudagraph_mode_across_dp,
    ) = get_batch_metadata_across_dp(
        num_tokens, cudagraph_size, cudagraph_runtime_mode, dp_size, dp_rank
    )
    if torch.all(num_tokens_across_dp == 0).item():
        # All ranks have zero tokens to run.
        return 0, None, 0

    # Synchronize cudagraph_runtime_mode across ranks by taking the minimum.
    synced_cudagraph_mode = int(cudagraph_mode_across_dp.min().item())
    # Check if all ranks have valid cudagraph_size.
    all_have_cudagraph = torch.all(cudagraph_size_across_dp != -1).item()

    if synced_cudagraph_mode != 0 and all_have_cudagraph:
        # All ranks use cudagraph. Pad to max cudagraph_size.
        max_cudagraph_size = int(cudagraph_size_across_dp.max().item())
        num_tokens_across_dp[:] = max_cudagraph_size
        return max_cudagraph_size, num_tokens_across_dp, synced_cudagraph_mode
    else:
        # Fall back to eager mode (no cudagraph).
        # Either some rank doesn't have cudagraph size or mode is NONE.
        synced_cudagraph_mode = 0
        num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
        num_tokens_after_padding = int(num_tokens_across_dp[dp_rank].item())
        return num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode


Review comment (performance):
```suggestion
    num_tokens_across_dp, cudagraph_size_across_dp, cudagraph_mode_across_dp = (
        get_batch_metadata_across_dp(
            num_tokens, cudagraph_size, cudagraph_runtime_mode, dp_size, dp_rank
        )
    )
```

Code after:


def get_cudagraph_and_dp_padding(
    num_tokens: int,
    cudagraph_size: int | None,
    cudagraph_runtime_mode: int,
    dp_size: int,
    dp_rank: int,
) -> tuple[int, torch.Tensor | None, int]:
    if dp_size == 1:
        if cudagraph_size is not None:
            return cudagraph_size, None, cudagraph_runtime_mode
        else:
            return num_tokens, None, cudagraph_runtime_mode

    # Convert None to -1 for sync (indicates no cudagraph available)
    if num_tokens == 0:
        cudagraph_size = 0
    elif cudagraph_size is None:
        cudagraph_size = -1

    num_tokens_across_dp, cudagraph_size_across_dp, cudagraph_mode_across_dp = (
        get_batch_metadata_across_dp(
            num_tokens, cudagraph_size, cudagraph_runtime_mode, dp_size, dp_rank
        )
    )
    if torch.all(num_tokens_across_dp == 0).item():
        # All ranks have zero tokens to run.
        return 0, None, 0

    # Synchronize cudagraph_runtime_mode across ranks by taking the minimum.
    synced_cudagraph_mode = int(cudagraph_mode_across_dp.min().item())
    # Check if all ranks have valid cudagraph_size.
    all_have_cudagraph = torch.all(cudagraph_size_across_dp != -1).item()

    if synced_cudagraph_mode != 0 and all_have_cudagraph:
        # All ranks use cudagraph. Pad to max cudagraph_size.
        max_cudagraph_size = int(cudagraph_size_across_dp.max().item())
        num_tokens_across_dp[:] = max_cudagraph_size
        return max_cudagraph_size, num_tokens_across_dp, synced_cudagraph_mode
    else:
        # Fall back to eager mode (no cudagraph).
        # Either some rank doesn't have cudagraph size or mode is NONE.
        synced_cudagraph_mode = 0
        num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
        num_tokens_after_padding = int(num_tokens_across_dp[dp_rank].item())
        return num_tokens_after_padding, num_tokens_across_dp, synced_cudagraph_mode