xai-org/x-algorithm · error · NotImplementedError

{q_seq_len=} must be a multiple of {config.block_q_dq=} * {c

Error message

{q_seq_len=} must be a multiple of {config.block_q_dq=} * {compute_wgs=}

What it means

The backward dkv/dq kernels split work across compute_wgs warp groups, so q_seq_len must be divisible by config.block_q_dq * compute_wgs (= compute_wgs_bwd). Remainder sequences cannot be tiled and raise NotImplementedError.

Source

Thrown at phoenix/xrex/pallas/ranker_attention_fa3.py:518

        raise NotImplementedError("Causal attention not supported in the backwards pass yet.")

    if not config.has_backward_blocks:
        raise ValueError("Need to specify backward blocks.")

    assert config.block_q_dq is not None
    assert config.block_kv_dq is not None
    assert config.block_q_dkv is not None
    assert config.block_kv_dkv is not None

    batch_size, q_seq_len, num_q_heads, head_dim = q.shape
    _, kv_seq_len, num_kv_heads, _ = k.shape
    q_heads_per_kv_head = num_q_heads // num_kv_heads
    dtype = q.dtype
    compute_wgs = config.compute_wgs_bwd

    num_q_tiles, rem = divmod(q_seq_len, config.block_q_dq * compute_wgs)
    if rem:
        raise NotImplementedError(
            f"{q_seq_len=} must be a multiple of {config.block_q_dq=} * {compute_wgs=}"
        )

    num_kv_tiles, rem = divmod(kv_seq_len, config.block_kv_dkv * compute_wgs)
    if rem:
        raise NotImplementedError(
            f"{kv_seq_len=} must be a multiple of {config.block_kv_dkv=} * {compute_wgs=}"
        )

    num_q_tiles_in_dkv, rem = divmod(q_seq_len, config.block_q_dkv)
    if rem:
        raise NotImplementedError(f"{q_seq_len=} must be a multiple of {config.block_q_dkv=}")

    num_kv_tiles_in_dq, rem = divmod(kv_seq_len, config.block_kv_dq)
    if rem:
        raise NotImplementedError(f"{kv_seq_len=} must be a multiple of {config.block_kv_dq=}")

    bound_arr = _normalize_bound(bound, batch_size, q_seq_len)

View on GitHub (pinned to 24c60942c5)

Solutions

Pad q (and kv consistently) to a multiple of block_q_dq * compute_wgs_bwd and mask with bound
Or pick block sizes whose product with compute_wgs divides q_seq_len (e.g. 64 divisors)

Example fix

# before
grads = jax.grad(loss)(params)  # q_seq_len=1000, block_q_dq=128, compute_wgs=2
# after
# pad q_seq_len to 1024 (multiple of 256) and mask with bound, or:
cfg = replace(cfg, block_q_dq=64)  # 64*2=128 divides 1000? no -> pad to 1024
q = jnp.pad(q, ((0,0),(0,24),(0,0),(0,0)))

Defensive patterns

Strategy: validation

Validate before calling

divisor = config.block_q_dq * config.compute_wgs_bwd
assert q_seq_len % divisor == 0, f"pad q_seq_len to a multiple of {divisor}"

Prevention

Compute backward divisors (block*_dq/dkv * compute_wgs_bwd) when choosing padding
Reuse one padded seq-len for forward and backward alignment checks

When it happens

Trigger: Differentiating with q_seq_len not a multiple of block_q_dq * compute_wgs_bwd, e.g. seq 1000 with block_q_dq=128 and compute_wgs=2 (divisor 256).

Common situations: Short fine-tuning sequences; adjusting compute_wgs_bwd or backward block sizes without re-checking sequence padding.

Related errors

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