xai-org/x-algorithm · error · ValueError

Block sparse tensors{context} {dim_name} dim must be {tgt} o

Error message

Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.

What it means

Validates that mask_block_cnt's batch and head dims match the expected count shape (or are 1 for broadcasting). The kernel only broadcasts size-1 dims; any other mismatch between the tensor's B/H and the model's batch/head count fails immediately.

Source

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

    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(
            f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}."
        )
    if expected_m_blocks != num_m_blocks:
        raise ValueError(
            f"Block sparse tensors{context} m-block dimension {num_m_blocks} does not match "
            f"sparse_block_size_q={sparse_block_size_q}. "
            f"Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE."
        )

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Set the mismatched dim to 1 so it broadcasts, or rebuild the mask with the correct batch/head dimensions
  2. Verify the num_head/batch arguments passed to normalize_block_sparse_config match the tensors' shapes
  3. If the mask is shared across layers with different head counts, expand it explicitly to the max head count

Example fix

// before
cnt = torch.load('mask.pt')           # (2, 32, M) but batch=4
cfg = normalize_block_sparse_config(tensors, batch_size=4, num_head=32, ...)

// after
cnt = torch.load('mask.pt')
if cnt.shape[0] not in (4, 1):
    cnt = cnt[:1].expand(4, -1, -1).contiguous()   # (4, 32, M)
cfg = normalize_block_sparse_config(tensors, batch_size=4, num_head=32, ...)
Defensive patterns

Strategy: validation

Validate before calling

def check_bh(t, dim, expected, allowed=(1,)):
    s = t.shape[dim]
    assert s == expected or s in allowed, f'dim {dim}: {s} != {expected}'
check_bh(tensors.mask_block_cnt, 0, batch); check_bh(tensors.mask_block_cnt, 1, num_head)

Prevention

When it happens

Trigger: normalize_block_sparse_config with mask_block_cnt whose shape[0] != batch (and != 1) or shape[1] != num_head (and != 1), e.g. running batch=4 inference with a mask built for batch=1 but reshaped to 2, or a per-layer head count differing from the mask's head dim.

Common situations: Reusing a cached/precomputed block mask across batch sizes or model variants (different num_attention_heads); H100 vs Blackwell configs with different head counts; forgetting that only literal 1 broadcasts.

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/218f56e7a72f9cd4. Report an issue: GitHub.