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 MHA backward pass dispatcher received a backward_pass_impl string it does not recognize. The implementation supports only specific named backward variants (the 'kv' variant shown above); anything else falls through to this ValueError.

Source

Thrown at phoenix/xrex/pallas/attention.py:702

                ),
            ],
            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)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Check the function signature/docstring of _mha_backward for the accepted backward_pass_impl values and use one of them (e.g. 'kv')
2. Fix typos in the config string
3. If you relied on a removed variant, pin the previous library version or migrate to the supported one

Example fix

# before
out = mha(q, k, v, ..., backward_pass_impl="kv_recompute")

# after
out = mha(q, k, v, ..., backward_pass_impl="kv")

VALID = {"kv"}
assert backward_pass_impl in VALID, f"backward_pass_impl must be one of {VALID}"

def is_valid_backward_impl(s: str) -> bool:
    return s in {"kv"}

When it happens

Trigger: Calling mha (or its vjp) with backward_pass_impl set to a value other than the supported implementation names (e.g. a typo like 'kv_cache' or a removed variant like 'dq').

Common situations: Upgrading versions where a backward variant was renamed/removed, typos in the attention config string, or copying example code from a different release.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• cap_method must be in [tanh, soft_sign], got {cap_method}
• cap_method must be in [tanh, soft_sign]
• Invalid backward pass implementation: {backward_pass_impl}
• Causal attention not supported in the backwards pass yet.
• Need to specify backward blocks.

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