xai-org/x-algorithm · error · ValueError

Only bfloat16 is supported for keys.

Error message

Only bfloat16 is supported for keys.

What it means

The top_k_by_key CUDA kernel (exposed via JAX FFI) only implements bfloat16 key comparisons; passing keys of any other dtype (float32, float16, int) is rejected upfront. Callers must cast scores/keys to jnp.bfloat16 before invoking top_k_by_key or its wrapper local_top_k.

Source

Thrown at phoenix/xrex/cuda/top_k_by_key/__init__.py:46

except ImportError:
    top_k_by_key_radix_select_api = None
else:
    jax.ffi.register_ffi_target(
        "xrex_top_k_by_key_radix_select",
        fn=top_k_by_key_radix_select_api.top_k_by_key_radix_select(),
        platform="CUDA",
    )


def top_k_by_key(
    keys: jax.Array,
    k: int,
    heuristic_pivot_ratio: float,
    use_async: bool = False,
    use_radix_select: bool = False,
):
    if keys.dtype != jnp.bfloat16:
        raise ValueError("Only bfloat16 is supported for keys.")
    if keys.ndim > 2:
        raise ValueError("keys must be 1D or 2D.")

    n = keys.shape[-1]
    if k > n:
        raise ValueError(f"k ({k}) must be <= n ({n})")

    if use_radix_select:
        api = top_k_by_key_radix_select_api
    elif use_async:
        api = top_k_by_key_async_api
    else:
        api = top_k_by_key_api
    if api is None or jax.default_backend() != "gpu":
        sorted_keys, sorted_indices = jax.lax.top_k(keys, k)
        return sorted_keys, sorted_indices.astype(jnp.int32)

    out_shape = (k,) if keys.ndim == 1 else (keys.shape[0], k)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Cast keys before the call: keys = keys.astype(jnp.bfloat16).
2. If bf16 rounding is unacceptable, fall back to jax.lax.top_k on the original dtype.
3. Keep the scoring head in bf16 end-to-end so no cast is needed at the top-k boundary.

Example fix

# before
idx, vals = top_k_by_key(scores, k=k, heuristic_pivot_ratio=0.5)  # scores is float32

# after
idx, vals = top_k_by_key(scores.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)

if keys.dtype != jnp.bfloat16:
    keys = keys.astype(jnp.bfloat16)
idx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)

def is_bf16_keys(keys: jax.Array) -> bool:
    return keys.dtype == jnp.bfloat16

try:
    idx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)
except ValueError as e:
    if "bfloat16" in str(e):
        idx, vals = top_k_by_key(keys.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)
    else:
        raise

When it happens

Trigger: Calling top_k_by_key(keys, k, ...) or local_top_k with keys.dtype in {float32, float16, int32, ...} — i.e. anything but jnp.bfloat16, including scores produced by a float32 dot-product or logit head.

Common situations: Feeding uncast model logits/similarity scores (commonly float32) into the candidate- retrieval top-k; refactoring a pipeline that previously used jax.lax.top_k (dtype-agnostic) to the fused CUDA kernel; mixed-precision training where activations are float32 at the scoring point.

Related errors

• keys must be 1D or 2D.
• {key}: got {arr.shape}/{arr.dtype}, manifest says {meta['sha
• k ({k}) must be <= n ({n})
• {name} must have dtype torch.int32
• Type mismatch: {self.q_dtype} != {self.k_dtype}

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