xai-org/x-algorithm · error · ValueError

Block sparse tensors

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.

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)
Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

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


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

Appendix: 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)