xai-org/x-algorithm · critical · ValueError

block-sparse arrays cover {bs_num_blocks} m-tiles but the ke

Error message

block-sparse arrays cover {bs_num_blocks} m-tiles but the kernel iterates {(seq_len + m_block - 1) // m_block} (seq_len={seq_len})

What it means

ranker_attention_fa4 derives the number of m-tiles it iterates from seq_len and m_block as ceil(seq_len / m_block), and requires the block-sparse layout's arrays (fwd_bs[3]) to cover exactly that many tiles. If the layout was built for a different sequence length or tile size, attention would read out-of-bounds or skip tiles, so it raises immediately.

Source

Thrown at phoenix/xrex/cutedsl/ranker_attention_fa4.py:100

    from xrex.cutedsl.ranker_fa4.flash_bwd_postprocess import FlashAttentionBackwardPostprocess
    from xrex.cutedsl.ranker_fa4.flash_bwd_sm100 import FlashAttentionBackwardSm100
    from xrex.cutedsl.ranker_fa4.flash_fwd_sm100 import FlashAttentionForwardSm100

    batch_size, seq_len, num_q_heads, head_dim = q.shape
    num_kv_heads = k.shape[2]
    qhead_per_kvhead = num_q_heads // num_kv_heads
    m_block = 128
    n_block = 128
    hdr = ((head_dim + 31) // 32) * 32
    sr_q = ((seq_len + m_block - 1) // m_block) * m_block
    sr_k = ((seq_len + n_block - 1) // n_block) * n_block
    dKV_postprocess = True

    fwd_bs, bwd_bs = block_sparse_layout
    bs_num_blocks = int(fwd_bs[3].shape[-2])
    bs_max_hist_blocks = int(fwd_bs[3].shape[-1])
    if bs_num_blocks != (seq_len + m_block - 1) // m_block:
        raise ValueError(
            f"block-sparse arrays cover {bs_num_blocks} m-tiles but the kernel "
            f"iterates {(seq_len + m_block - 1) // m_block} (seq_len={seq_len})"
        )
    use_pack_gqa = qhead_per_kvhead > 1 and (m_block % qhead_per_kvhead == 0)

    cache_key = (
        head_dim,
        num_q_heads,
        num_kv_heads,
        batch_size,
        seq_len,
        bs_num_blocks,
        bs_max_hist_blocks,
        use_pack_gqa,
    )

    if cache_key not in _FA4_KERNEL_CACHE:
        fa_fwd = FlashAttentionForwardSm100(

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Rebuild the block-sparse layout with the same seq_len and m_block used for the kernel call
  2. Align padding: ensure seq_len passed to the kernel equals the seq_len used to construct fwd_bs
  3. If layouts are cached, key the cache on (seq_len, m_block) so stale layouts are not reused

Example fix

// before
layout = build_block_sparse_layout(seq_len=1024, ...)
out = ranker_attention_fa4(..., seq_len=2048, m_block=128, block_sparse_layout=layout)
// after
layout = build_block_sparse_layout(seq_len=2048, ...)
out = ranker_attention_fa4(..., seq_len=2048, m_block=128, block_sparse_layout=layout)
Defensive patterns

Strategy: validation

Validate before calling

bs_num_blocks = int(fwd_bs[3].shape[-2])
expected = (seq_len + m_block - 1) // m_block
assert bs_num_blocks == expected, f"stale layout: {bs_num_blocks} vs {expected}"

Try / catch

try:
    out = ranker_attention_fa4(...)
except ValueError as e:
    if "m-tiles" in str(e):
        layout = build_block_sparse_layout(seq_len=seq_len, m_block=m_block, ...)
        out = ranker_attention_fa4(...)
    else:
        raise

Prevention

When it happens

Trigger: Calling ranker_attention_fa4 (usually via sharded_mha) with a block_sparse_layout built with a different seq_len or m_block than the one passed to the kernel; mismatch between padded/unpadded sequence lengths.

Common situations: Changing seq_len (e.g. different batch padding or packing) without rebuilding build_block_sparse_layout; altering m_block/block size config in one place but not the other; caching layouts across configs.

Related errors


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