xai-org/x-algorithm · error · ValueError

cap_method must be in [tanh, soft_sign], got {cap_method}

Error message

cap_method must be in [tanh, soft_sign], got {cap_method}

What it means

The varlen ranker attention forward kernel applies logit capping when cap > 0.0 and only implements 'tanh' and 'soft_sign' methods; any other cap_method string triggers this guard inside the scan body. It mirrors the v2 kernel's validation and indicates a bad hyperparameter reaching the kernel.

Source

Thrown at phoenix/xrex/pallas/ranker_attention_varlen.py:282

        )
        seg_k = pl.load(
            segment_ref,
            (pl.dslice(offset_k, block_k),),
            mask=k_mask,
            other=PADDING_SEGMENT_ID,
        ).astype(jnp.int8)[None, :]

        qk = pl.dot(q, k, trans_b=True)
        if sm_scale != 1.0:
            qk *= sm_scale

        if cap > 0.0:
            if cap_method == "tanh":
                qk = cap * tanh(qk / cap)
            elif cap_method == "soft_sign":
                qk = qk / (1.0 + jnp.abs(qk) / cap)
            else:
                raise ValueError(f"cap_method must be in [tanh, soft_sign], got {cap_method}")

        mask = jnp.logical_or(seg_k > 0, span_q[:, None] == span_k[None, :])
        mask = jnp.logical_and(seq_q_is_not_padding, mask)
        mask = jnp.logical_and(mask, k_mask[None, :])

        if causal:
            causal_mask = span_q[:, None] >= span_k[None, :]
            mask = jnp.logical_and(mask, causal_mask)

        qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
        max_logit = jnp.max(qk, axis=1)
        max_logit = jnp.maximum(max_logit, 1.0)
        m_curr = jnp.maximum(m_prev, max_logit)
        p = jnp.exp(qk - m_curr[:, None])
        alpha = jnp.exp(m_prev - m_curr)
        l_prev = l_prev * alpha + jnp.sum(p, axis=1)

        p = p.astype(q.dtype)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Use cap_method='tanh' or cap_method='soft_sign'.
2. Disable capping with cap=0.0 if that was the intent.
3. Validate cap_method in config parsing before kernel launch.

Example fix

# before
attn_varlen(..., cap=50.0, cap_method='sinh')
# after
attn_varlen(..., cap=50.0, cap_method='tanh')

assert cap <= 0.0 or cap_method in ('tanh', 'soft_sign')

from typing import Literal
CapMethod = Literal['tanh', 'soft_sign']
def is_valid_cap_method(s: str) -> bool:
    return s in ('tanh', 'soft_sign')

When it happens

Trigger: Calling varlen ranker attention with cap > 0.0 and cap_method not in ['tanh', 'soft_sign'].

Common situations: Hyperparameter sweeps with unsupported capping names; configs ported from FlashAttention-style APIs that use different cap method names; typos/case errors.

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

• Invalid backward pass implementation: {backward_pass_impl}
• Invalid backward pass implementation: {backward_pass_impl}
• 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.
• {head_dim=} must be divisible by 64

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