{"record":{"id":"444a2e212ea92af6","repo":"xai-org/x-algorithm","slug":"only-bfloat16-is-supported-for-keys","errorCode":null,"errorMessage":"Only bfloat16 is supported for keys.","messagePattern":"Only bfloat16 is supported for keys\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/top_k_by_key/__init__.py","lineNumber":46,"sourceCode":"except ImportError:\n    top_k_by_key_radix_select_api = None\nelse:\n    jax.ffi.register_ffi_target(\n        \"xrex_top_k_by_key_radix_select\",\n        fn=top_k_by_key_radix_select_api.top_k_by_key_radix_select(),\n        platform=\"CUDA\",\n    )\n\n\ndef top_k_by_key(\n    keys: jax.Array,\n    k: int,\n    heuristic_pivot_ratio: float,\n    use_async: bool = False,\n    use_radix_select: bool = False,\n):\n    if keys.dtype != jnp.bfloat16:\n        raise ValueError(\"Only bfloat16 is supported for keys.\")\n    if keys.ndim > 2:\n        raise ValueError(\"keys must be 1D or 2D.\")\n\n    n = keys.shape[-1]\n    if k > n:\n        raise ValueError(f\"k ({k}) must be <= n ({n})\")\n\n    if use_radix_select:\n        api = top_k_by_key_radix_select_api\n    elif use_async:\n        api = top_k_by_key_async_api\n    else:\n        api = top_k_by_key_api\n    if api is None or jax.default_backend() != \"gpu\":\n        sorted_keys, sorted_indices = jax.lax.top_k(keys, k)\n        return sorted_keys, sorted_indices.astype(jnp.int32)\n\n    out_shape = (k,) if keys.ndim == 1 else (keys.shape[0], k)","sourceCodeStart":28,"sourceCodeEnd":64,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/top_k_by_key/__init__.py#L28-L64","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast keys before the call: keys = keys.astype(jnp.bfloat16).","If bf16 rounding is unacceptable, fall back to jax.lax.top_k on the original dtype.","Keep the scoring head in bf16 end-to-end so no cast is needed at the top-k boundary."],"exampleFix":"# before\nidx, vals = top_k_by_key(scores, k=k, heuristic_pivot_ratio=0.5)  # scores is float32\n\n# after\nidx, vals = top_k_by_key(scores.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)","handlingStrategy":"type-guard","validationCode":"if keys.dtype != jnp.bfloat16:\n    keys = keys.astype(jnp.bfloat16)\nidx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)","typeGuard":"def is_bf16_keys(keys: jax.Array) -> bool:\n    return keys.dtype == jnp.bfloat16","tryCatchPattern":"try:\n    idx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)\nexcept ValueError as e:\n    if \"bfloat16\" in str(e):\n        idx, vals = top_k_by_key(keys.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)\n    else:\n        raise","preventionTips":["Keep the scoring head in bf16 so no boundary cast is needed.","Wrap the kernel call in a small helper that normalizes dtype and rank once."],"tags":["cuda-kernel","dtype","bfloat16","top-k","retrieval"],"backgroundTag":"unsupported-dtype-for-kernel","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}