vllm.models.step5.amd.checks

Startup checks for the ROCm Step5 sparse attention layers.

Functions:

• validate_parallel_geometry –

  The indexer scores regions per provider group and attention reads the

validate_parallel_geometry(*, total_num_heads, total_num_kv_heads, indexer_num_heads, index_tp_size, tp_size)

The indexer scores regions per provider group and attention reads the matching KV group, so both partitions must agree globally and per rank.

Source code in vllm/models/step5/amd/checks.py

def validate_parallel_geometry(
    *,
    total_num_heads: int,
    total_num_kv_heads: int,
    indexer_num_heads: int,
    index_tp_size: int,
    tp_size: int,
) -> None:
    """The indexer scores regions per provider group and attention reads the
    matching KV group, so both partitions must agree globally and per rank."""
    if index_tp_size != total_num_kv_heads:
        raise ValueError(
            "Step5 sparse provider groups must match attention KV groups: "
            f"{index_tp_size} != {total_num_kv_heads}"
        )
    if indexer_num_heads % index_tp_size:
        raise ValueError(
            f"{indexer_num_heads} sparse indexer heads do not divide into "
            f"{index_tp_size} provider groups."
        )
    if total_num_heads % tp_size:
        raise ValueError(
            f"Step5 attention heads ({total_num_heads}) must be divisible by "
            f"the tensor parallel size ({tp_size})."
        )
    local_kv_groups = _groups_per_rank(total_num_kv_heads, tp_size)
    q_heads_per_kv, rem = divmod(total_num_heads // tp_size, local_kv_groups)
    if rem or q_heads_per_kv not in (4, 8, 16):
        raise ValueError(
            "Step5 sparse attention needs 4, 8 or 16 local query heads per KV "
            f"group (num_heads={total_num_heads}, tp_size={tp_size})."
        )