Skip to content

vllm.models.common.ops.sequence_parallel

Functions:

can_sp_all_gather_matmul(x, weight)

Whether sp_all_gather_matmul applies to this shard and weight.

Source code in vllm/models/common/ops/sequence_parallel.py 
60
61
62
63
64
65
66
67
68
def can_sp_all_gather_matmul(x: torch.Tensor, weight: torch.Tensor) -> bool:
    """Whether ``sp_all_gather_matmul`` applies to this shard and weight."""
    return (
        x.shape[0] >= _FUSED_AG_MATMUL_MIN_TOKENS_PER_RANK
        and x.dtype == weight.dtype == torch.bfloat16
        and weight.ndim == 2
        and not torch.cuda.is_current_stream_capturing()
        and _has_fused_all_gather_matmul()
    )

is_sequence_parallel_active(is_enabled, num_tokens, min_tokens=0) cached

Whether a batch of num_tokens shards the stream along the sequence.

Below one token per rank the shards are mostly padding, so the batch stays replicated. A model may raise this lower bound with min_tokens.

Source code in vllm/models/common/ops/sequence_parallel.py 
25
26
27
28
29
30
31
32
33
34
35
36
@lru_cache
def is_sequence_parallel_active(
    is_enabled: bool, num_tokens: int, min_tokens: int = 0
) -> bool:
    """Whether a batch of ``num_tokens`` shards the stream along the sequence.

    Below one token per rank the shards are mostly padding, so the batch
    stays replicated. A model may raise this lower bound with ``min_tokens``.
    """
    return is_enabled and num_tokens >= max(
        get_tensor_model_parallel_world_size(), min_tokens
    )

sp_all_gather_matmul(x, weight)

sp_all_gather(x) @ weight.T, overlapping the gather with the GEMM.

Source code in vllm/models/common/ops/sequence_parallel.py 
71
72
73
74
75
76
77
78
79
def sp_all_gather_matmul(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """``sp_all_gather(x) @ weight.T``, overlapping the gather with the GEMM."""
    _, outputs = torch.ops.symm_mem.fused_all_gather_matmul(
        x.contiguous(),
        [weight.t()],
        gather_dim=0,
        group_name=get_tp_group().device_group.group_name,
    )
    return outputs[0]