{"record":{"id":"3734ab5ebb491b8c","repo":"jax-ml/jax","slug":"require-q-seqlen-and-kv-seqlen-to-use-packed-layou","errorCode":null,"errorMessage":"Require q_seqlen and kv_seqlen to use packed layout","messagePattern":"Require q_seqlen and kv_seqlen to use packed layout","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":2142,"sourceCode":"          f\"Expected 'None' for bias, mask, q_seqlen, and kv_seqlen, \"\n          f\"but got: bias={bias}, mask={mask}, q_seqlen={q_seqlen}, kv_seqlen={kv_seqlen}\"\n      )\n    check_fp8_params(fp8_params)\n    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets,\n      None, None, layout)\n    output, amax_s, amax_o = _dot_product_attention_fp8(\n        query, key, value, fp8_params,\n        scale, mask_type == MaskType.CAUSAL, layout.value, cudnn_version\n    )\n    return output, amax_s, amax_o\n  else:\n    if has_padding(mask_type) and (q_seqlen is None or kv_seqlen is None):\n        raise ValueError(\"Require q_seqlen and kv_seqlen to generate padding mask\")\n    if sliding_window_length is not None and sliding_window_length <= 0:\n      raise ValueError(\n        f\"Require sliding_window_length > 0, got {sliding_window_length}\")\n    if q_offsets is not None and (q_seqlen is None or kv_seqlen is None):\n      raise ValueError(\"Require q_seqlen and kv_seqlen to use packed layout\")\n\n    # A bias gradient can only be needed if a differentiable operand feeds the\n    # combined bias: an explicit bias, or a non-boolean mask. A boolean mask\n    # alone is converted to a constant additive bias whose gradient nobody can\n    # request, so skip the (costly) dbias computation in the backward pass.\n    bias_is_differentiable = bias is not None or (\n        mask is not None and mask.dtype != np.dtype('bool'))\n    bias = combine_bias_and_mask(bias, mask, query.dtype)\n    # check if input shape and data type is compatiable\n    check_layout(query, key, value, bias, q_seqlen, kv_seqlen, q_offsets, kv_offsets,\n      None, None, layout)\n    has_bias = bias is not None\n    has_dbias = has_bias and bias_is_differentiable and \\\n      should_export_dbias(bias.shape, query.shape, layout)\n    variadic_args = (has_bias, has_dbias)\n\n    _not_used = jnp.zeros(0, dtype=query.dtype)\n    if bias is None:","sourceCodeStart":2124,"sourceCodeEnd":2160,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L2124-L2160","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pass q_seqlen and kv_seqlen arrays alongside q_offsets (both must be non-None for packed layout)","If you only have offsets, compute cumulative sequence lengths from them before calling","If you intended a padded layout instead, drop q_offsets and pass a mask_type with q_seqlen/kv_seqlen"],"exampleFix":"// before\nout = jax.nn.dot_product_attention(q, k, v, q_offsets=q_offsets)\n// after\nout = jax.nn.dot_product_attention(q, k, v, q_offsets=q_offsets,\n                                    q_seqlen=q_seqlen, kv_seqlen=kv_seqlen)","handlingStrategy":"validation","validationCode":"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'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always build (q_offsets, q_seqlen, kv_seqlen) together in varlen data loaders","Add a unit test covering the varlen path"],"tags":["jax","attention","varlen","packed-layout","cudnn"],"backgroundTag":"missing-required-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}