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

Raised in _mha_backward when the backward_pass_impl argument is not one of the supported implementation names. The code dispatches between pallas kernels (e.g. mha_backward_kv) based on a string selector, and any unrecognized string falls into the else branch.

Source

Thrown at phoenix/xrex/pallas/ranker_attention.py:781

                ),
            ],
            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=2),
        )(q, k, v, temp, segment_ids, 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)


def mha_reference(
    q,
    k,
    v,
    sm_scale=1.0,
    cap=-1.0,
    cap_method="tanh",
    inverted_sliding_window_sizep1: int = 0,
):
    logits = jnp.einsum("bqhc,bkhc->bhqk", q, k).astype(jnp.float32)
    logits *= sm_scale
    if cap > 0.0:

View on GitHub (pinned to 24c60942c5)

Solutions

1. Check the valid backward_pass_impl values in _mha_backward's if/elif chain (phoenix/xrex/pallas/ranker_attention.py around line 770-781)
2. Pass one of the supported implementation names, e.g. 'kv'
3. If you need a different impl, upgrade/downgrade to the version that supports it

Example fix

// before
out = mha(q, k, v, backward_pass_impl="kv_split")
// after
out = mha(q, k, v, backward_pass_impl="kv")

valid_bwd = {"kv"}  # extend from the if/elif chain in _mha_backward
assert backward_pass_impl in valid_bwd, f"unsupported {backward_pass_impl=}"

null

When it happens

Trigger: Calling mha (with defvjp) so JAX triggers the VJP, while passing a backward_pass_impl string other than the supported kernel names (e.g. a typo like 'kv_nosplit' or 'flash').

Common situations: Typos in config, or code written against an older/newer version where the set of backward implementations changed; also passing None or empty string explicitly.

Related errors

• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign], got {cap_method}
• cap_method must be in [tanh, soft_sign]
• cap_method must be in [tanh, soft_sign, none], got {cap_meth
• q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.

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