xai-org/x-algorithm · error · ValueError

valid_block_upper and valid_block_lower must be provided tog

Error message

valid_block_upper and valid_block_lower must be provided together

What it means

ranker_attention_varlen_fa4 takes optional valid_block_upper and valid_block_lower masks that define per-block validity. They are coupled: passing exactly one of them is almost certainly a bug, so the function raises unless both are provided (or both are None, in which case zero-filled defaults are created from the layout shape).

Source

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

    from xrex.cutedsl.ranker_fa4.block_sparsity import BlockSparseTensors
    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, packed_S, num_q_heads, head_dim = q.shape
    num_kv_heads = k.shape[2]
    qpk = num_q_heads // num_kv_heads
    block_size = 128
    hdr = ((head_dim + 31) // 32) * 32
    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",

View on GitHub (pinned to 24c60942c5)

Solutions

1. Pass both valid_block_upper and valid_block_lower together
2. Or pass neither to accept the zero-filled defaults
3. Audit call sites in sharded_mha for one-sided mask construction

Example fix

# before
out = ranker_attention_varlen_fa4(..., valid_block_upper=ub)  # lower missing
# after
out = ranker_attention_varlen_fa4(..., valid_block_upper=ub, valid_block_lower=lb)

assert (valid_block_upper is None) == (valid_block_lower is None), \
    "valid_block_upper/lower must be passed together"

When it happens

Trigger: Calling ranker_attention_varlen_fa4 (typically via sharded_mha) with valid_block_upper set but valid_block_lower None (or vice versa).

Common situations: Copy-paste or partial refactor where one mask variable is renamed or dropped; conditionally computing only the upper bound for a causal-ish mask and forgetting the lower bound defaults to None.

Related errors

• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• k ({k}) must be <= n ({n})
• block-sparse arrays cover {bs_num_blocks} m-tiles but the ke
• Unsupported tcgen05 MMA op kind: {type(op).__name__}
• type checking expression %s failed: invalid argument type: e

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