jax-ml/jax · error · ValueError

cuDNN doesn't support right window: {r_window} when causal m

Error message

cuDNN doesn't support right window: {r_window} when causal mask is not used.

What it means

With implementation='cudnn', a sliding window (local_window_size) with a nonzero right window is only supported for causal masks. cuDNN's sliding-window attention requires the window to be one-sided (left-only) unless causal masking is used, so r_window != 0 with non-causal mask_type raises this.

Source

Thrown at jax/_src/nn/functions.py:1256

        if key_value_seq_lengths is None:
          key_value_seq_lengths = jnp.full((B,), S, dtype=np.int32)

      mask_type = MaskType.NO_MASK
      if use_padding and is_causal:
        mask_type = MaskType.PADDING_CAUSAL
      elif is_causal:
        mask_type = MaskType.CAUSAL
      elif use_padding:
        mask_type = MaskType.PADDING
      # CuDNN supports only the left window with an exclusive boundary when
      # causal mask is enabled.
      sliding_window = None
      if local_window_size is not None:
        l_window, r_window = local_window_size
        if r_window == 0 or mask_type == MaskType.CAUSAL:
          sliding_window = l_window + 1
        else:
          raise ValueError(f"cuDNN doesn't support right window: {r_window} "
                           "when causal mask is not used.")

      out = cudnn_dot_product_attention(
          query_arr, key_arr, value_arr, bias, mask, query_seq_lengths,
          key_value_seq_lengths, scale=scale_val, mask_type=mask_type,
          sliding_window_length=sliding_window, return_residual=return_residual,
      )
      if return_residual:
        # Regardless of input layout, cudnn always returns residual with
        # (B N T) layout.
        out, residual = out
        residual = jnp.transpose(residual, (0, 2, 1)).astype(out.dtype)
        out = (out, residual)
    case None:
      # TODO(kaixih@nvidia) Automatically select the best backend (defaults to XLA for now).
      out = _dot_product_attention_xla(
          query_arr, key_arr, value_arr, bias, mask, is_causal=is_causal,
          scale=scale_val, q_seqlen=query_seq_lengths,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a left-only window: local_window_size=(l, 0), or (l, 1) is not allowed—must be 0 for non-causal; set r_window=0
If you need a right window, use implementation='xla' (or default) instead of 'cudnn'
Set mask_type=MaskType.CAUSAL if the layer is causal, which enables l_window+1 sliding windows

Example fix

// before
jax.nn.dot_product_attention(q, k, v, implementation='cudnn',
    local_window_size=(128, 128))  # right window, not causal

// after
jax.nn.dot_product_attention(q, k, v, implementation='cudnn',
    local_window_size=(128, 0))  # left-only window

Defensive patterns

Strategy: fallback

Validate before calling

def supports_cudnn_window(l_window, r_window, mask_type):
    return r_window == 0 or mask_type == 'causal'  # MaskType.CAUSAL
if not supports_cudnn_window(l, r, mask_type):
    impl = 'xla'  # fall back instead of 'cudnn'

Prevention

Default to left-only windows (l, 0) for local attention
Select implementation='cudnn' only after checking mask/window compatibility

When it happens

Trigger: Calling dot_product_attention(..., implementation='cudnn', local_window_size=(l, r)) with r > 0 and mask_type not MaskType.CAUSAL (e.g. MaskType.NO_MASK or additive).

Common situations: Porting local/banded attention (e.g. from a model using bilateral windows) to the cuDNN flash path; forgetting to pass mask_type=MaskType.CAUSAL for decoder layers.

Related errors

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