xai-org/x-algorithm · error · ValueError

Block sparse tensors{context} n-block dimension must be <= {

Error message

Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}.

What it means

The last dimension of mask_block_idx (the max number of N/KV blocks indexed per row) must not exceed expected_n_blocks, the total available KV blocks. Indices pointing past the end of the KV axis would be out of bounds in the kernel.

Source

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

        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."
        )
    return expected_count_shape, expected_index_shape, q_subtile_factor


def get_block_sparse_expected_shapes_bwd(
    batch_size: int,
    num_head: int,
    seqlen_q: int,
    seqlen_k: int,
    m_block_size: int,
    n_block_size: int,

View on GitHub (pinned to 24c60942c5)

Solutions

1. Regenerate the block mask for the current seqlen_k and sparse_block_size_kv so n-block dim <= ceil(seqlen_k / block_kv)
2. Clip or rebuild indices so the widest row fits within expected_n_blocks
3. Check that expected_n_blocks is computed with the same block size as the mask generator used

Example fix

# before
idx = build_mask(seqlen_k=8192)   # n dim = 8192/512 = 16
cfg = normalize_block_sparse_config(..., seqlen_k=4096)  # expects <= 8

# after
idx = build_mask(seqlen_k=4096)   # n dim = 8
cfg = normalize_block_sparse_config(..., seqlen_k=4096)

import math
exp_n = math.ceil(seqlen_k / sparse_block_kv)
assert tensors.mask_block_idx.shape[3] <= exp_n, tensors.mask_block_idx.shape

When it happens

Trigger: Passing mask_block_idx with shape[3] > expected_n_blocks, e.g. a mask built for a longer KV sequence (larger seqlen_k or smaller sparse_block_size_kv) than the current run.

Common situations: Reusing a mask generated for a longer context window; changing seqlen_k or the KV block size without regenerating the mask; off-by-one in n-block computation in a custom mask builder.

Related errors

• {name}_block_cnt and {name}_block_idx must be on the same de
• {name}_block tensors must live on CUDA
• {name} must have dtype torch.int32
• {name} must be on the same device as block sparse tensors
• {name} must live on CUDA

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