jax-ml/jax · error · ValueError

{soft_cap=} must not be 0.0.

Error message

{soft_cap=} must not be 0.0.

What it means

In ragged_paged_attention, soft_cap (tanh attention logit capping, as used by Gemma-2) is optional, but exactly 0.0 is forbidden because dividing by a zero cap would blow up the logits; use None to disable capping. static_validate_inputs distinguishes None from 0.0 explicitly.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/kernel.py:274

        f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if (
      kv_lens.dtype != jnp.int32
      or page_indices.dtype != jnp.int32
      or cu_q_lens.dtype != jnp.int32
  ):
    raise ValueError(
        "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
        f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
        f" {cu_q_lens.dtype=}."
    )
  if num_q_heads % num_kv_heads != 0:
    raise ValueError(f"{num_q_heads=} must be divisible by {num_kv_heads=}")
  if sliding_window is not None and sliding_window <= 0:
    raise ValueError(f"{sliding_window=} must be positive.")
  if soft_cap is not None and soft_cap == 0.0:
    raise ValueError(f"{soft_cap=} must not be 0.0.")
  if (
      num_kv_pages_per_block is not None
      and not 0 < num_kv_pages_per_block <= pages_per_seq
  ):
    raise ValueError(
        f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]."
    )
  if num_queries_per_block is not None and num_queries_per_block <= 0:
    raise ValueError(f"{num_queries_per_block=} must be positive.")
  if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:
    raise ValueError(f"{vmem_limit_bytes=} must be positive.")
  del sm_scale  # No constraints on sm_scale.
  del mask_value  # No consstraints on mask_value.


def ragged_paged_attention_kernel(
    # Prefetch
    kv_lens_ref,  # [max_num_seqs]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert 0.0 to None before calling: soft_cap = soft_cap or None (careful: this also maps a genuine tiny cap; check == 0.0 explicitly)
2. Pass the model's real cap value (e.g. 50.0 for Gemma-2) when capping is desired
3. Normalize model configs at load time so 'disabled' is always None

Example fix

// before
soft_cap = cfg.attn_logit_softcapping  # 0.0 when disabled
attn(..., soft_cap=soft_cap)
// after
soft_cap = None if not cfg.attn_logit_softcapping else cfg.attn_logit_softcapping
attn(..., soft_cap=soft_cap)

soft_cap = None if soft_cap is None or soft_cap == 0.0 else float(soft_cap)

When it happens

Trigger: Passing soft_cap=0.0 intending to turn off logit soft-capping.

Common situations: Config-driven code where a missing attn_logit_softcapping value defaults to 0.0 (common in Gemma-2 / HF configs) and is forwarded verbatim to the TPU kernel.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Number of cores or threads must be at least 1, but got {self
• Vector clock size must be at least 1, but got {self.vector_c
• {num_q_heads=} must be divisible by {num_kv_heads=}
• {sliding_window=} must be positive.
• Not implemented: {num_combined_kv_heads=} can not be XLA ful

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0f95236d3f35f4b5. Report an issue: GitHub.