jax-ml/jax · error · ValueError

Require q_seqlen and kv_seqlen to use packed layout

Error message

Require q_seqlen and kv_seqlen to use packed layout

What it means

Raised by the cuDNN fused attention StableHLO path in jax when the caller supplies q_offsets (indicating a packed/varlen query layout) but does not also supply q_seqlen and kv_seqlen tensors. The packed layout needs per-sequence lengths to reconstruct sequence boundaries, so the call is rejected before tracing.

Source

Thrown at jax/_src/cudnn/fused_attention_stablehlo.py:2142

          f"Expected 'None' for bias, mask, q_seqlen, and kv_seqlen, "
          f"but got: bias={bias}, mask={mask}, q_seqlen={q_seqlen}, kv_seqlen={kv_seqlen}"
      )
    check_fp8_params(fp8_params)
    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets,
      None, None, layout)
    output, amax_s, amax_o = _dot_product_attention_fp8(
        query, key, value, fp8_params,
        scale, mask_type == MaskType.CAUSAL, layout.value, cudnn_version
    )
    return output, amax_s, amax_o
  else:
    if has_padding(mask_type) and (q_seqlen is None or kv_seqlen is None):
        raise ValueError("Require q_seqlen and kv_seqlen to generate padding mask")
    if sliding_window_length is not None and sliding_window_length <= 0:
      raise ValueError(
        f"Require sliding_window_length > 0, got {sliding_window_length}")
    if q_offsets is not None and (q_seqlen is None or kv_seqlen is None):
      raise ValueError("Require q_seqlen and kv_seqlen to use packed layout")

    # A bias gradient can only be needed if a differentiable operand feeds the
    # combined bias: an explicit bias, or a non-boolean mask. A boolean mask
    # alone is converted to a constant additive bias whose gradient nobody can
    # request, so skip the (costly) dbias computation in the backward pass.
    bias_is_differentiable = bias is not None or (
        mask is not None and mask.dtype != np.dtype('bool'))
    bias = combine_bias_and_mask(bias, mask, query.dtype)
    # check if input shape and data type is compatiable
    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets,
      None, None, layout)
    has_bias = bias is not None
    has_dbias = has_bias and bias_is_differentiable and \
      should_export_dbias(bias.shape, query.shape, layout)
    variadic_args = (has_bias, has_dbias)

    _not_used = jnp.zeros(0, dtype=query.dtype)
    if bias is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass q_seqlen and kv_seqlen arrays alongside q_offsets (both must be non-None for packed layout)
2. If you only have offsets, compute cumulative sequence lengths from them before calling
3. If you intended a padded layout instead, drop q_offsets and pass a mask_type with q_seqlen/kv_seqlen

Example fix

// before
out = jax.nn.dot_product_attention(q, k, v, q_offsets=q_offsets)
// after
out = jax.nn.dot_product_attention(q, k, v, q_offsets=q_offsets,
                                    q_seqlen=q_seqlen, kv_seqlen=kv_seqlen)

assert q_offsets is None or (q_seqlen is not None and kv_seqlen is not None), 'packed layout requires q_seqlen and kv_seqlen'

When it happens

Trigger: Calling jax.nn.dot_product_attention (or the cudnn fused attention path) with q_offsets set but q_seqlen=None or kv_seqlen=None, e.g. building a varlen batch from ragged sequences.

Common situations: Migrating from a padding-mask API to the packed layout while forgetting the seqlen arrays; passing q_offsets together with a padding mask_type without the corresponding lengths; assuming offsets alone encode sequence lengths.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• cuDNN doesn't support right window: {r_window} when causal m
• Unsupported sequence length Q {T}, KV {S}.
• Packed layout requires a GPU with at least Hopper architectu
• mla requires cudnn version >= 9.10 and at least hopper arch.
• cuDNN is not detected.

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