xai-org/x-algorithm · error · ValueError

Block sparse tensors{context} m-block dimension {num_m_block

Error message

Block sparse tensors{context} m-block dimension {num_m_blocks} does not match sparse_block_size_q={sparse_block_size_q}. Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE.

What it means

The number of M blocks implied by the mask metadata (num_m_blocks, derived from expected_count_shape / seqlen_q) does not equal expected_m_blocks computed from sparse_block_size_q. The BlockSparseTensorsTorch.block_size (BlockMask BLOCK_SIZE) must match the query-block size the kernel uses.

Source

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

        ("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,
    subtile_factor: int,
) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]:
    sparse_block_size_q = subtile_factor * m_block_size
    expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q)

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Set BlockSparseTensorsTorch.block_size (and the BlockMask BLOCK_SIZE) so that seqlen_q / block_size[0] equals the kernel's expected m-block count
  2. Regenerate the mask with the kernel's sparse_block_size_q (see normalize_block_sparse_config docs)
  3. Print sparse_block_size_q and num_m_blocks from the error context and align the mask generator to them

Example fix

# before
tensors = BlockSparseTensorsTorch(..., block_size=(128, 128))  # M blocks = seqlen/128

# after
tensors = BlockSparseTensorsTorch(..., block_size=(sparse_block_size_q, tile_n))  # matches kernel
# i.e. rebuild BlockMask with BLOCK_SIZE matching the ranker FA4 config
Defensive patterns

Strategy: validation

Validate before calling

assert tensors.block_size is None or tensors.block_size[0] * num_m_blocks == seqlen_q, \
    (tensors.block_size, num_m_blocks, seqlen_q)

Prevention

When it happens

Trigger: Setting BlockSparseTensorsTorch.block_size[0] (or the BlockMask BLOCK_SIZE) to a value inconsistent with seqlen_q / the kernel's m-block layout; e.g. block_size=(128, 128) when the kernel derives 64-row m-blocks from seqlen_q.

Common situations: Porting a FlexAttention BlockMask created with one BLOCK_SIZE into this ranker FA4 path that assumes a specific sparse_block_size_q; mixing forward block sizes into a config without recomputing m-block count.

Related errors


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