xai-org/x-algorithm · error · ValueError

Invalid backward pass implementation: {backward_pass_impl}

Error message

Invalid backward pass implementation: {backward_pass_impl}

What it means

The backward pass of the fused MHA kernel dispatches on backward_pass_impl, selecting a specialized Pallas kernel per implementation (e.g. the 'kv' fused-dkv kernel). If the string does not match any known implementation, this ValueError is raised. It is a top-level API argument validation failure for _mha_backward.

Source

Thrown at phoenix/xrex/pallas/ranker_attention_v2.py:1091

                ),
            ],
            out_specs=[
                pl.BlockSpec(
                    index_map=lambda j, k, _: (j, 0, k, 0),
                    block_shape=(None, seq_len, None, head_dim),
                ),
                pl.BlockSpec(
                    index_map=lambda j, k, _: (j, 0, k, 0),
                    block_shape=(None, seq_len, None, head_dim),
                ),
            ],
            name="mha_backward_kv",
            debug=debug,
            interpret=interpret,
            compiler_params=CompilerParams(num_warps=num_warps, num_stages=num_stages),
        )(q, k, v, temp, segment_ids, k_block_bwd_dkv_state_metadata, out, do_scaled, l, m, delta)
    else:
        raise ValueError(f"Invalid backward pass implementation: {backward_pass_impl}")
    return dq.astype(q.dtype), dk, dv, dtemp, dsegment


mha.defvjp(_mha_forward, _mha_backward)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Inspect the function signature/default of _mha_backward in ranker_attention_v2.py to list valid backward_pass_impl values and use one of them.
2. Omit backward_pass_impl to use the default implementation.
3. Pin or align the library version with the config you copied the value from.

Example fix

# before
out = mha(q, k, v, backward_pass_impl='flash')
# after
out = mha(q, k, v)  # or backward_pass_impl='kv' if explicitly needed

import inspect
sig = inspect.signature(_mha_backward)  # or the public wrapper
# omit backward_pass_impl to use the default, or verify against
# the branches in phoenix/xrex/pallas/ranker_attention_v2.py

try:
    out = mha(q, k, v, backward_pass_impl=impl)
except ValueError as e:
    if 'Invalid backward pass implementation' in str(e):
        out = mha(q, k, v)  # fall back to default
    else:
        raise

When it happens

Trigger: Calling the ranker attention v2 public API with backward_pass_impl set to a value other than the supported ones (e.g. 'kv' / supported impl names), such as 'flash', 'vanilla', a typo, or an option removed/renamed between versions.

Common situations: Upgrading the library and passing a backward_pass_impl name that no longer exists; copying example configs from incompatible versions; experiment sweeps enumerating unsupported values.

Related errors

• Causal attention not supported in the backwards pass yet.
• Need to specify backward blocks.
• {q_seq_len=} must be a multiple of {config.block_q_dq=} * {c
• cap_method must be in [tanh, soft_sign], got {cap_method}
• cap_method must be in [tanh, soft_sign]

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