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 will iterate {_expected_m_blocks} (packed_S={packed_S}). Pass packed_seq_len (the physical packed row length) to build_block_sparse_layout so every physical tile has an entry.

What it means

The varlen FA4 kernel iterates ceil(packed_S / block_size) physical m-tiles over the packed sequence rows, and every physical tile must have an entry in the block-sparse layout's index array (bs_num_blocks from fwd_bs[3]). If the layout was built from logical per-sequence lengths instead of the physical packed row length, some tiles would have no sparsity entry and the kernel would fault or misbehave, so it raises with guidance to pass packed_seq_len to build_block_sparse_layout.

Source

Thrown at phoenix/xrex/cutedsl/ranker_attention_varlen_fa4.py:337

    sr_q = ((packed_S + block_size - 1) // block_size) * block_size
    sr_k = sr_q
    dKV_postprocess = True
    use_pack_gqa = qpk > 1 and (block_size % qpk == 0)

    fwd_bs, _ = block_sparse_layout
    if valid_block_upper is None or valid_block_lower is None:
        if valid_block_upper is not None or valid_block_lower is not None:
            raise ValueError("valid_block_upper and valid_block_lower must be provided together")
        valid_block_upper = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32)
        valid_block_lower = jnp.zeros(fwd_bs[2].shape, dtype=jnp.int32)
    valid_block_upper = jnp.broadcast_to(valid_block_upper, fwd_bs[2].shape)
    valid_block_lower = jnp.broadcast_to(valid_block_lower, fwd_bs[2].shape)
    bs_max_hist_blocks = int(fwd_bs[3].shape[-1])
    bs_num_blocks = int(fwd_bs[3].shape[-2])

    _expected_m_blocks = (packed_S + block_size - 1) // block_size
    if bs_num_blocks != _expected_m_blocks:
        raise ValueError(
            f"block-sparse arrays cover {bs_num_blocks} m-tiles but the kernel "
            f"will iterate {_expected_m_blocks} (packed_S={packed_S}). "
            "Pass packed_seq_len (the physical packed row length) to "
            "build_block_sparse_layout so every physical tile has an entry."
        )

    cache_key = (
        "packed",
        head_dim,
        num_q_heads,
        num_kv_heads,
        batch_size,
        packed_S,
        bs_num_blocks,
        bs_max_hist_blocks,
        use_pack_gqa,
    )

View on GitHub (pinned to 24c60942c5)

Solutions

Pass packed_seq_len (the physical packed row length, matching packed_S given to the kernel) to build_block_sparse_layout
Ensure packed_S and the layout are derived from the same cu_seqlens/packing buffer
Add an assertion in your pipeline: build_block_sparse_layout(packed_seq_len=packed_S, ...) right where packed_S is computed so they cannot drift

Example fix

# before
layout = build_block_sparse_layout(seq_len=max_len, ...)
out = ranker_attention_varlen_fa4(..., packed_S=packed_S, block_size=bs, block_sparse_layout=layout)
# after
layout = build_block_sparse_layout(packed_seq_len=packed_S, ...)
out = ranker_attention_varlen_fa4(..., packed_S=packed_S, block_size=bs, block_sparse_layout=layout)

Defensive patterns

Strategy: validation

Validate before calling

expected_m_blocks = (packed_S + block_size - 1) // block_size
assert int(fwd_bs[3].shape[-2]) == expected_m_blocks, \
    "layout built without packed_seq_len; rebuild with packed_seq_len=packed_S"

Prevention

Compute packed_S once and thread it to both build_block_sparse_layout and the kernel
Add a helper that builds layout + calls the kernel so lengths cannot diverge
Property-test packing with ragged batch sizes

When it happens

Trigger: Calling ranker_attention_varlen_fa4 via sharded_mha where build_block_sparse_layout received total/max logical seq_len (sum of cu_seqlens minus padding, or max len) instead of packed_S = physical number of packed rows.

Common situations: Migrating from the non-varlen FA4 path where seq_len was the right argument; padding the packed buffer to a multiple and passing the padded logical length; computing packed_S from batch size x max_len instead of the actual packed row count.

Related errors

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