xai-org/x-algorithm · error · ValueError

Block sparse tensors{context} must have shapes (B, H, M) and

Error message

Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N).

What it means

Raised by infer_block_sparse_expected_shapes when validating block-sparse attention mask tensors. The count tensor mask_block_cnt must be 3-D (B, H, M) and the index tensor mask_block_idx must be 4-D (B, H, M, N); any other rank is rejected before shapes are compared. This is an early structural check so downstream broadcast/expansion logic can assume fixed ranks.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/block_sparsity.py:343

    if sparse_block_size_q % base_m_block != 0:
        raise ValueError(
            f"Block sparse tensors{context} have block size {sparse_block_size_q}, "
            f"which must be a multiple of {base_m_block}."
        )

    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)
    expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv)
    q_subtile_factor = sparse_block_size_q // base_m_block
    expected_count_shape = (batch_size, num_head, expected_m_blocks)
    expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks)

    mask_block_cnt = tensors.mask_block_cnt
    mask_block_idx = tensors.mask_block_idx
    if mask_block_cnt is None or mask_block_idx is None:
        raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.")
    if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4:
        raise ValueError(
            f"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N)."
        )
    for dim_name, cur, tgt in (
        ("batch", mask_block_cnt.shape[0], expected_count_shape[0]),
        ("head", mask_block_cnt.shape[1], expected_count_shape[1]),
    ):
        if cur != tgt and cur != 1:
            raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
    for dim_name, cur, tgt in (
        ("batch", mask_block_idx.shape[0], expected_index_shape[0]),
        ("head", mask_block_idx.shape[1], expected_index_shape[1]),
    ):
        if cur != tgt and cur != 1:
            raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.")
    if mask_block_cnt.shape[2] != mask_block_idx.shape[2]:
        raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.")
    if mask_block_idx.shape[3] > expected_n_blocks:
        raise ValueError(

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Reshape mask_block_cnt to (B, H, M) and mask_block_idx to (B, H, M, N) before passing them in
  2. If you built these from a BlockMask, use the repo's provided conversion utilities instead of hand-rolling tensor indexing
  3. Print .shape of both tensors right before the call and compare against the message's expected shapes

Example fix

// before
cnt = counts.squeeze()        # now (H, M)
idx = indices                 # (B, H, N)

// after
cnt = counts.squeeze(0).unsqueeze(0) if counts.ndim == 2 else counts  # (B, H, M)
idx = indices.unsqueeze(2) if indices.ndim == 3 else indices          # (B, H, M, N)
Defensive patterns

Strategy: validation

Validate before calling

def check_block_sparse_ranks(tensors):
    assert tensors.mask_block_cnt is not None and tensors.mask_block_idx is not None
    assert tensors.mask_block_cnt.ndim == 3, tensors.mask_block_cnt.shape
    assert tensors.mask_block_idx.ndim == 4, tensors.mask_block_idx.shape

Type guard

def has_valid_block_ranks(t: BlockSparseTensorsTorch) -> bool:
    return (t.mask_block_cnt is not None and t.mask_block_idx is not None
            and t.mask_block_cnt.ndim == 3 and t.mask_block_idx.ndim == 4)

Prevention

When it happens

Trigger: Calling normalize_block_sparse_config (directly or via the ranker FA4 forward path) with BlockSparseTensorsTorch whose mask_block_cnt is not ndim==3 or mask_block_idx is not ndim==4, e.g. passing a 2-D count tensor or a 5-D index tensor built from a custom BlockMask conversion.

Common situations: Converting a FlexAttention/BlockMask seqlens or indices tensor without adding the batch/head dims; accidentally stacking or squeezing a dim; mismatch between varlen (2-D) layouts and the dense (B, H, ...) layout this kernel expects.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/eb12db26e85ddfdf. Report an issue: GitHub.