jax-ml/jax · error · ValueError

{num_seqs.shape=} must be (1,)

Error message

{num_seqs.shape=} must be (1,)

What it means

static_validate_inputs runs at trace/compile time on the ragged paged-attention path (used by ragged_paged_attention, ref impl, and dynamic validation). num_seqs must be a shape-(1,) device array holding the live sequence count; other shapes (scalar, (n,), etc.) are rejected because the kernel reads num_seqs[0] as a single value.

Source

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

    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = None,
    k_scale: float | None = None,
    v_scale: float | None = None,
    # Kernel tuning params.
    num_kv_pages_per_block: int | None = None,
    num_queries_per_block: int | None = None,
    vmem_limit_bytes: int | None = None,
):
  _, num_q_heads, head_dim = q.shape
  _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape
  assert num_combined_kv_heads % 2 == 0
  assert isinstance(k_scale, float) or k_scale is None
  assert isinstance(v_scale, float) or v_scale is None
  num_kv_heads = num_combined_kv_heads // 2
  max_num_seqs, pages_per_seq = page_indices.shape
  if num_seqs.shape != (1,):
    raise ValueError(f"{num_seqs.shape=} must be (1,)")
  if head_dim_k != head_dim:
    raise ValueError(
        f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}."
    )
  if kv_lens.shape != (max_num_seqs,):
    raise ValueError(
        f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where"
        " `max_num_seqs` is `page_indices.shape[0]`."
    )
  if cu_q_lens.shape != (max_num_seqs + 1,):
    raise ValueError(
        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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the count: num_seqs = jnp.array([n], dtype=jnp.int32)
2. Keep scheduler output shapes stable: num_seqs always shape (1,)

Example fix

// before
ragged_paged_attention(q, k_pages, page_indices, num_seqs=5, ...)
// after
ragged_paged_attention(q, k_pages, page_indices, num_seqs=jnp.array([5], jnp.int32), ...)

num_seqs = jnp.asarray(num_seqs).reshape(1)

def is_num_seqs_shape(a): return isinstance(a, jax.Array) and a.shape == (1,)

When it happens

Trigger: Passing num_seqs as a Python int, a 0-d array, or a per-sequence array of counts instead of jnp.array([n], dtype=int). Note this must be an array even though the count is dynamic.

Common situations: Passing num_seqs=5 (int) for convenience; reshaping scheduler state; migrating from an API that accepted a scalar count.

Related errors

• Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},) whe
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.
• group_offset must be a ()-shaped array. Got: {group_offset.s
• k_pages and v_pages must have the same shape. Got {k_pages.s

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