vllm.v1.attention.backends.flash_attn

Attention layer with FlashAttention.

Classes:

• FA4DenseAttentionKernel –

  Warm paged, causal FA4 on Blackwell and Hopper.

• FlashAttentionBackend –
• FlashAttentionImpl –
• FlashAttentionMetadataBuilder –

Functions:

• use_cascade_attention –

  Decide whether to use cascade attention.

FA4DenseAttentionKernel

Bases: VllmJitKernel['FA4DenseAttentionKernel.CompileKey']

Warm paged, causal FA4 on Blackwell and Hopper.

Source code in vllm/v1/attention/backends/flash_attn.py

class FA4DenseAttentionKernel(VllmJitKernel["FA4DenseAttentionKernel.CompileKey"]):
    """Warm paged, causal FA4 on Blackwell and Hopper."""

    @dataclass(frozen=True)
    class CompileKey:
        q_stage: int
        dtype: torch.dtype
        qhead_per_kvhead: int
        head_dim: int
        page_size: int
        has_window_left: bool
        has_window_right: bool
        softcap: float
        is_single_batch: bool
        num_splits: int

    @staticmethod
    def kernel(*args: Any, **kwargs: Any) -> Any:
        assert flash_attn_varlen_func is not None
        return flash_attn_varlen_func(*args, **kwargs)

    def dispatch(  # type: ignore[override]
        self,
        *,
        q_stage: int,
        dtype: torch.dtype,
        num_qo_heads: int,
        num_kv_heads: int,
        head_dim: int,
        page_size: int,
        window_size: tuple[int, int],
        softcap: float,
        is_single_batch: bool = True,
        num_splits: int = 1,
    ) -> CompileKey:
        return self.CompileKey(
            q_stage=q_stage,
            dtype=dtype,
            qhead_per_kvhead=num_qo_heads // num_kv_heads,
            head_dim=head_dim,
            page_size=page_size,
            has_window_left=window_size[0] >= 0,
            has_window_right=window_size[1] >= 0,
            softcap=softcap,
            is_single_batch=is_single_batch,
            num_splits=num_splits,
        )

    def get_warmup_keys(
        self,
        *,
        vllm_config: VllmConfig,
        kv_cache_spec: AttentionSpec,
        num_qo_heads: int,
    ) -> list[CompileKey]:
        capability = current_platform.get_device_capability()
        major = capability.major if capability is not None else None
        if (
            vllm_config.parallel_config.decode_context_parallel_size != 1
            or vllm_config.model_config.rswa_window is not None
            or kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
            or kv_cache_spec.dtype not in FA4_DENSE_FLOAT_DTYPES
            or not isinstance(kv_cache_spec, (FullAttentionSpec, SlidingWindowSpec))
        ):
            return []
        head_dim = kv_cache_spec.head_size
        page_size = kv_cache_spec.block_size
        sliding_window = kv_cache_spec.sliding_window
        window_size = (
            (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
        )
        softcap = float(
            getattr(
                vllm_config.model_config.hf_text_config, "attn_logit_softcapping", 0
            )
            or 0
        )
        if (
            major in (10, 11)
            and uses_fa4_hd256_kernel(head_dim)
            and page_size == FA4_HD256_PAGE_SIZE
        ):
            q_stages = WarmupChoices(1, 2)
            single_batches = WarmupChoices(True, False)
            split_counts = WarmupChoices(1)
        elif (
            major == 9 and head_dim == 512 and window_size == (-1, -1) and softcap == 0
        ):
            # SM90 forward has split/non-split variants; its transitive
            # combine specializes at 32/64/128/256 splits.
            q_stages = WarmupChoices(1)
            single_batches = WarmupChoices(False)
            split_counts = WarmupChoices(1, 32, 64, 128, 256)
        else:
            return []

        return self._trace_dispatch(self.dispatch)(
            q_stage=q_stages,
            is_single_batch=single_batches,
            num_splits=split_counts,
            dtype=vllm_config.model_config.dtype,
            num_qo_heads=num_qo_heads,
            num_kv_heads=kv_cache_spec.num_kv_heads,
            head_dim=head_dim,
            page_size=page_size,
            window_size=WarmupChoices(window_size),
            softcap=softcap,
        )

    def compile(self, compile_key: CompileKey) -> None:
        assert compile_flash_attn_varlen_func_from_specs is not None
        max_seqlen_q = FA4_DENSE_Q_TILE + 1 if compile_key.q_stage == 2 else 1
        batch_size = 1 if compile_key.is_single_batch else 2
        kv_shape = (
            FA4_DENSE_NUM_BLOCKS,
            compile_key.page_size,
            1,
            compile_key.head_dim,
        )
        kv_stride = (
            2 * compile_key.page_size * compile_key.head_dim,
            2 * compile_key.head_dim,
            2 * compile_key.head_dim,
            1,
        )
        compile_flash_attn_varlen_func_from_specs(
            q_shape=(
                batch_size * max_seqlen_q,
                compile_key.qhead_per_kvhead,
                compile_key.head_dim,
            ),
            k_shape=kv_shape,
            v_shape=kv_shape,
            k_stride=kv_stride,
            v_stride=kv_stride,
            q_dtype=compile_key.dtype,
            cu_seqlens_q_shape=(batch_size + 1,),
            seqused_k_shape=(batch_size,),
            page_table_shape=(
                batch_size,
                FA4_DENSE_MAX_SEQLEN_K // compile_key.page_size,
            ),
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=FA4_DENSE_MAX_SEQLEN_K,
            causal=True,
            window_size=[
                1 if compile_key.has_window_left else -1,
                1 if compile_key.has_window_right else -1,
            ],
            softcap=compile_key.softcap,
            num_splits=2 if compile_key.num_splits > 1 else 1,
            fa_version=4,
            return_softmax_lse=False,
        )
        if compile_key.head_dim == 512 and compile_key.num_splits > 1:
            from torch._subclasses.fake_tensor import FakeTensorMode

            from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd_combine

            # FA4's forward compile-only entry returns before its transitive
            # SplitKV combine. Call the native combine entry in FakeTensorMode:
            # it populates the native cache without launching the kernel.
            with FakeTensorMode():
                _flash_attn_fwd_combine(
                    torch.empty(
                        (compile_key.num_splits, 2, 1, 512),
                        dtype=torch.float32,
                        device="cuda",
                    ),
                    torch.empty(
                        (compile_key.num_splits, 2, 1),
                        dtype=torch.float32,
                        device="cuda",
                    ),
                    torch.empty((2, 1, 512), dtype=compile_key.dtype, device="cuda"),
                    cu_seqlens=torch.empty((3,), dtype=torch.int32, device="cuda"),
                )

    def __call__(
        self,
        *,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        **kwargs: Any,
    ) -> Any:
        return self.kernel(q=q, k=k, v=v, **kwargs)

FlashAttentionBackend

Bases: AttentionBackend

Methods:

• supports_attn_type –

  FlashAttention supports all attention types.

Source code in vllm/v1/attention/backends/flash_attn.py

class FlashAttentionBackend(AttentionBackend):
    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
        "auto",
        "float16",
        "bfloat16",
        "fp8",
        "fp8_e4m3",
    ]
    head_size_v: int | None = None

    @staticmethod
    def _get_sm90_fa4_fp8_kv_block_size(
        kv_cache_spec: KVCacheSpec | None = None,
    ) -> int | None:
        if kv_cache_spec is None:
            vllm_config = get_current_vllm_config_or_none()
            if vllm_config is None or vllm_config.model_config is None:
                return None
            head_size = vllm_config.model_config.get_head_size()
            uses_fp8_kv = vllm_config.cache_config.cache_dtype in (
                "fp8",
                "fp8_e4m3",
            )
        else:
            if not isinstance(kv_cache_spec, AttentionSpec):
                return None
            head_size = kv_cache_spec.head_size
            uses_fp8_kv = kv_cache_spec.kv_quant_mode == KVQuantMode.FP8_PER_TENSOR

        if (
            current_platform.is_device_capability_family(90)
            and uses_fp8_kv
            and head_size == 512
            and get_flash_attn_version(head_size=head_size) == 4
        ):
            # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page.
            return 64
        return None

    @classmethod
    def _get_fa4_hd256_block_size(
        cls, kv_cache_spec: KVCacheSpec | None = None
    ) -> int | None:
        if kv_cache_spec is None:
            vllm_config = get_current_vllm_config_or_none()
            if vllm_config is None or vllm_config.model_config is None:
                return None
            head_size = vllm_config.model_config.get_head_size()
            head_size_v = cls.head_size_v
        else:
            if not isinstance(kv_cache_spec, AttentionSpec):
                return None
            head_size = kv_cache_spec.head_size
            head_size_v = kv_cache_spec.head_size_v

        if (
            uses_fa4_hd256_kernel(head_size, head_size_v)
            and get_flash_attn_version(
                head_size=head_size,
                head_size_v=head_size_v,
                supports_fa4_hd256=True,
            )
            == 4
        ):
            return FA4_HD256_PAGE_SIZE
        return None

    @classmethod
    def get_supported_kernel_block_sizes(
        cls, kv_cache_spec: KVCacheSpec | None = None
    ) -> list[int | MultipleOf]:
        if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec):
            return [block_size]
        if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec):
            # Sliding-window specs select the smallest advertised size.
            return [block_size]
        return [MultipleOf(16)]

    forward_includes_kv_cache_update: bool = False

    @classmethod
    def get_preferred_block_size(cls, default_block_size: int) -> int:
        if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
            return max(default_block_size, block_size)
        if block_size := cls._get_fa4_hd256_block_size():
            return max(default_block_size, block_size)
        if current_platform.is_xpu():
            return max(default_block_size, 64)
        return super().get_preferred_block_size(default_block_size)

    @staticmethod
    def get_name() -> str:
        return "FLASH_ATTN"

    @classmethod
    def supports_rswa(cls) -> bool:
        return True

    @classmethod
    def supports_sliding_window(cls) -> bool:
        return True

    @classmethod
    def supports_batch_invariance(cls) -> bool:
        return True

    @classmethod
    def supports_non_causal(cls) -> bool:
        return True

    @classmethod
    def supports_attn_type(cls, attn_type: str) -> bool:
        """FlashAttention supports all attention types."""
        return attn_type in (
            AttentionType.DECODER,
            AttentionType.ENCODER,
            AttentionType.ENCODER_ONLY,
            AttentionType.ENCODER_DECODER,
        )

    @classmethod
    def supports_per_head_quant_scales(cls) -> bool:
        fa_version = get_flash_attn_version()
        return fa_version is not None and fa_version >= 3

    @staticmethod
    def get_impl_cls() -> type["FlashAttentionImpl"]:
        return FlashAttentionImpl

    @staticmethod
    def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
        return FlashAttentionMetadataBuilder

    @classmethod
    def supports_head_size(cls, head_size: int) -> bool:
        if head_size % 8 != 0:
            return False
        if head_size <= 256:
            return True
        if is_fa_version_supported(4):
            return head_size <= 512
        return False

    @classmethod
    def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
        if kv_cache_dtype is None:
            return True
        if kv_cache_dtype not in cls.supported_kv_cache_dtypes:
            return False
        if is_quantized_kv_cache(kv_cache_dtype):
            return flash_attn_supports_kv_cache_dtype(kv_cache_dtype)
        return True

    @classmethod
    def supports_mm_prefix(cls) -> bool:
        return is_fa_version_supported(4)

    @classmethod
    def supports_sink(cls) -> bool:
        if not is_flash_attn_varlen_func_available():
            return False
        return flash_attn_supports_sinks()

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability >= DeviceCapability(8, 0)

    @classmethod
    def supports_combination(
        cls,
        head_size: int,
        dtype: torch.dtype,
        kv_cache_dtype: CacheDType | None,
        block_size: int | None,
        use_mla: bool,
        has_sink: bool,
        use_sparse: bool,
        use_mm_prefix: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
        if has_sink and device_capability < DeviceCapability(9, 0):
            return "sink not supported on compute capability < 9.0"
        if (
            use_mm_prefix
            and kv_cache_dtype is not None
            and is_quantized_kv_cache(kv_cache_dtype)
            and device_capability == DeviceCapability(9, 0)
        ):
            return "SM90 FP8 KV with mm_prefix requires Triton"
        if (
            kv_cache_dtype is not None
            and is_quantized_kv_cache(kv_cache_dtype)
            and not flash_attn_supports_kv_cache_dtype(
                kv_cache_dtype,
                head_size=head_size,
                head_size_v=head_size,
                has_sinks=has_sink,
                kv_cache_block_size=block_size,
                supports_fa4_hd256=True,
            )
        ):
            return (
                "FP8 KV cache requires FA3 on SM90, FA4 with head_size=512 "
                "on SM90, or FA4 on SM100"
            )
        if (
            use_mm_prefix
            and get_flash_attn_version(
                head_size=head_size,
                has_sinks=has_sink,
                kv_cache_block_size=block_size,
                supports_fa4_hd256=True,
            )
            != 4
        ):
            return (
                "mm_prefix (PrefixLM bidirectional attention) requires "
                "FlashAttention v4, which does not resolve for this "
                "head_size"
            )
        return None

supports_attn_type(attn_type) classmethod

FlashAttention supports all attention types.

Source code in vllm/v1/attention/backends/flash_attn.py

@classmethod
def supports_attn_type(cls, attn_type: str) -> bool:
    """FlashAttention supports all attention types."""
    return attn_type in (
        AttentionType.DECODER,
        AttentionType.ENCODER,
        AttentionType.ENCODER_ONLY,
        AttentionType.ENCODER_DECODER,
    )

FlashAttentionImpl

Bases: AttentionImpl

Methods:

• forward –

  Forward pass with FlashAttention.

Source code in vllm/v1/attention/backends/flash_attn.py

class FlashAttentionImpl(AttentionImpl):
    can_return_lse_for_decode: bool = True
    supports_dcp: bool = True

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: list[float] | None,
        sliding_window: int | None,
        kv_cache_dtype: str,
        logits_soft_cap: float | None = None,
        attn_type: AttentionType = AttentionType.DECODER,
        kv_sharing_target_layer_name: str | None = None,
        sinks: torch.Tensor | None = None,
    ) -> None:
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        if alibi_slopes is not None:
            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
        self.alibi_slopes = alibi_slopes
        if sliding_window is None:
            self.sliding_window = (-1, -1)
        elif attn_type == AttentionType.ENCODER_ONLY:
            self.sliding_window = (sliding_window - 1, sliding_window - 1)
        else:
            self.sliding_window = (sliding_window - 1, 0)
        self.kv_cache_dtype = kv_cache_dtype
        if logits_soft_cap is None:
            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
            logits_soft_cap = 0
        self.logits_soft_cap = logits_soft_cap
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

        self.attn_type = attn_type
        vllm_config = get_current_vllm_config_or_none()
        uses_kv_cache = attn_type not in (
            AttentionType.ENCODER,
            AttentionType.ENCODER_ONLY,
        )
        # The final KV cache block size is unavailable during construction.
        self.vllm_flash_attn_version = get_flash_attn_version(
            requires_alibi=alibi_slopes is not None,
            head_size=head_size,
            has_sinks=sinks is not None,
            requires_softcap=bool(self.logits_soft_cap),
            supports_fa4_hd256=True,
        )
        self.fa4_hd256 = self.vllm_flash_attn_version == 4 and uses_fa4_hd256_kernel(
            head_size
        )
        if self.fa4_hd256 and not uses_kv_cache and sliding_window is not None:
            # The hd256 kernel requires seqused_k for local attention.
            logger.warning_once(
                "FA4's Blackwell head_size=256 kernel does not support local "
                "attention on encoder inputs, defaulting to FA version 2."
            )
            self.vllm_flash_attn_version = 2
            self.fa4_hd256 = False
        logger.info_once(
            "Using FlashAttention version %s",
            self.vllm_flash_attn_version,
        )
        # Cache the batch invariant result for use in forward passes
        self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT

        if is_quantized_kv_cache(
            self.kv_cache_dtype
        ) and not flash_attn_supports_kv_cache_dtype(
            self.kv_cache_dtype,
            requires_alibi=alibi_slopes is not None,
            head_size=head_size,
            head_size_v=head_size,
            has_sinks=sinks is not None,
            requires_softcap=bool(self.logits_soft_cap),
            supports_fa4_hd256=True,
        ):
            raise NotImplementedError(
                f"FlashAttention does not support {self.kv_cache_dtype}"
                " kv-cache on this device."
            )

        self.sinks = sinks
        if self.sinks is not None:
            assert flash_attn_supports_sinks(), (
                "Sinks are only supported in FlashAttention 3"
            )
            assert self.sinks.shape[0] == num_heads, (
                "Sinks must have the same number of heads as the number of "
                "heads in the layer"
            )

        # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes
        # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q,
        # K, and V to have the same FP8 dtype.
        uses_sm90_fa4_fp8_kv_dequant = (
            self.vllm_flash_attn_version == 4
            and current_platform.is_device_capability_family(90)
            and self.kv_cache_dtype in ("fp8", "fp8_e4m3")
        )
        self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (
            not uses_sm90_fa4_fp8_kv_dequant
        )

        dcp_a2a = (
            vllm_config is not None
            and vllm_config.parallel_config.decode_context_parallel_size > 1
            and vllm_config.parallel_config.dcp_comm_backend == "a2a"
        )
        self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs

        self._dcp_dtype: torch.dtype | None = None
        self._dcp_max_num_tokens: int = 0
        if vllm_config is not None and self.dcp_world_size > 1:
            self._dcp_dtype = vllm_config.model_config.dtype
            self._dcp_max_num_tokens = (
                vllm_config.scheduler_config.max_num_batched_tokens
            )

    def forward(
        self,
        layer: torch.nn.Module,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        output: torch.Tensor,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass with FlashAttention.

        Args:
            layer: The attention layer, providing the q/k/v quantization scales.
            query: shape = [num_tokens, num_heads, head_size]
            key: shape = [num_tokens, num_kv_heads, head_size]
            value: shape = [num_tokens, num_kv_heads, head_size]
            kv_cache: shape =
                [num_blocks, num_kv_heads, block_size, 2 * head_size]
            attn_metadata: Metadata for attention.
            output: Tensor that the attention result is written into.
            output_scale: Scale for fused output quantization; not supported
                by this backend.
            output_block_scale: Block scale for fused output quantization;
                not supported by this backend.

        Returns:
            shape = [num_tokens, num_heads * head_size]
        NOTE: FP8 quantization, flash-attn expect the size of
              {q,k,v}_descale to be (num_sequences, num_kv_heads).
              We use torch's .expand() to avoid duplicating values

        """
        assert self.vllm_flash_attn_version is not None, (
            "FlashAttention version not detected."
        )

        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError(
                "fused output quantization is not yet supported for FlashAttentionImpl"
            )

        if attn_metadata is None:
            # Profiling run.
            return output.fill_(0)

        attn_type = self.attn_type

        # IMPORTANT!
        # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
        # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
        # in this method. For example, `view` and `slice` (or `[:n]`) operations
        # are surprisingly slow even in the case they do not invoke any GPU ops.
        # Minimize the PyTorch ops in this method as much as possible.
        # Whenever making a change in this method, please benchmark the
        # performance to make sure it does not introduce any overhead.

        num_actual_tokens = attn_metadata.num_actual_tokens

        # Handle encoder attention differently - no KV cache needed
        if attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
            # For encoder attention,
            # we use direct Q, K, V tensors without caching
            return self._forward_encoder_attention(
                query[:num_actual_tokens],
                key[:num_actual_tokens],
                value[:num_actual_tokens],
                output[:num_actual_tokens],
                attn_metadata,
                layer,
            )

        # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
        # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
        # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
        # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
        fixed_k = canonicalize_singleton_dim_strides(key_cache)
        fixed_v = canonicalize_singleton_dim_strides(value_cache)
        if fixed_k is not key_cache or fixed_v is not value_cache:
            logger.debug(
                "Canonicalized degenerate KV cache strides (FlashAttention): "
                "shape=%s, key strides before=%s after=%s, "
                "value strides before=%s after=%s",
                key_cache.shape,
                key_cache.stride(),
                fixed_k.stride(),
                value_cache.stride(),
                fixed_v.stride(),
            )
        key_cache, value_cache = fixed_k, fixed_v

        if is_quantized_kv_cache(self.kv_cache_dtype):
            # queries are quantized in the attention layer
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())

        if not attn_metadata.use_cascade:
            cu_seqlens_q = attn_metadata.query_start_loc
            seqused_k = attn_metadata.seq_lens
            max_seqlen_q = attn_metadata.max_query_len
            max_seqlen_k = attn_metadata.max_seq_len
            block_table = attn_metadata.block_table
            scheduler_metadata = attn_metadata.scheduler_metadata

            descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)

            q_descale = (
                layer._q_scale.expand(descale_shape)
                if self.supports_quant_query_input
                else None
            )
            k_descale = layer._k_scale.expand(descale_shape)
            v_descale = layer._v_scale.expand(descale_shape)

            if self.dcp_world_size > 1:
                self._forward_with_dcp(
                    query[:num_actual_tokens],
                    key[:num_actual_tokens],
                    value[:num_actual_tokens],
                    key_cache,
                    value_cache,
                    output[:num_actual_tokens],
                    attn_metadata,
                    q_descale=q_descale,
                    k_descale=k_descale,
                    v_descale=v_descale,
                )
                return output
            else:
                causal = attn_metadata.causal
                is_dynamic_causal = isinstance(causal, torch.Tensor)

                # The layer's own window wins over the group's: one KV cache
                # group can hold both windowed and global layers (e.g. Gemma-3
                # with the hybrid KV cache manager disabled), and the group spec
                # cannot describe both.
                window = _maybe_symmetrize_window(self.sliding_window, causal)
                sliding_window_size: list[int] | None = (
                    list(window) if window is not None else None
                )

                mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
                mm_mask_mod = None
                mm_aux = None
                if (
                    mm_prefix_query_ranges is not None
                    and not is_dynamic_causal
                    and causal is True
                    and self.vllm_flash_attn_version == 4
                ):
                    # Triton convention: 1 + window_size[0]. Global layers store
                    # (-1, -1) → sw stays None.
                    layer_window = self.sliding_window
                    sw_val = (
                        1 + layer_window[0]
                        if layer_window is not None and layer_window[0] >= 0
                        else None
                    )
                    # Gemma4: also clamp the bidirectional block to the
                    # sliding window when the layer opts in
                    # (mm_prefix_clamp_sliding_window flag from PR #47217).
                    mm_clamp_sw = 0
                    if (
                        getattr(layer, "mm_prefix_clamp_sliding_window", False)
                        and sw_val is not None
                    ):
                        mm_clamp_sw = sw_val
                    mm_mask_mod = _make_mm_prefix_mask_mod(
                        sliding_window=mm_clamp_sw,
                        sliding_window_left=sw_val,
                    )
                    mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc]
                    # mm_prefix is (causal ∧ window) ∨ bidirectional-range —
                    # not ⊆ causal. FA #155 stopped auto-clearing causal/local
                    # when mask_mod is set, so the caller must disable them or
                    # the built-in causal path shorts out / clips the mask_mod.
                    causal = False
                    sliding_window_size = None

                # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
                # mask without block-size approximation.  The mask_mod encodes
                # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which
                # supersedes any FA-layer sliding_window_size parameter.
                rswa_mask_mod_fn = None
                rswa_aux = None
                if (
                    attn_metadata.rswa_prefix_lens is not None
                    and self.vllm_flash_attn_version == 4
                    and not is_dynamic_causal
                ):
                    rswa_mask_mod_fn = _make_rswa_mask_mod()
                    rswa_aux = [
                        attn_metadata.rswa_prefix_lens.to(torch.int32),
                        attn_metadata.rswa_window_tensor,  # pre-allocated CUDA tensor
                    ]
                    # mask_mod fully expresses R-SWA; disable FA's own window.
                    sliding_window_size = None

                dynamic_causal = None
                if isinstance(causal, torch.Tensor):
                    if self.vllm_flash_attn_version != 4:
                        raise NotImplementedError(
                            "Per-sequence causal requires FA4. Current version: "
                            f"FA{self.vllm_flash_attn_version}"
                        )
                    dynamic_causal = causal
                    has_window = (
                        sliding_window_size is not None and sliding_window_size[1] >= 0
                    )
                    causal = not has_window

                num_splits = attn_metadata.max_num_splits
                if self.fa4_hd256:
                    # hd256 requires page-aligned lengths, exact-width block
                    # tables, and no SplitKV.
                    num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
                    max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
                    block_table = block_table[:, :num_pages]
                    num_splits = 1

                _FA4_DENSE_ATTENTION_KERNEL(
                    q=query[:num_actual_tokens],
                    k=key_cache,
                    v=value_cache,
                    out=output[:num_actual_tokens],
                    cu_seqlens_q=cu_seqlens_q,
                    max_seqlen_q=max_seqlen_q,
                    seqused_k=seqused_k,
                    max_seqlen_k=max_seqlen_k,
                    softmax_scale=self.scale,
                    causal=causal,
                    alibi_slopes=self.alibi_slopes,
                    window_size=sliding_window_size,
                    block_table=block_table,
                    softcap=self.logits_soft_cap,
                    scheduler_metadata=scheduler_metadata,
                    fa_version=self.vllm_flash_attn_version,
                    q_descale=q_descale,
                    k_descale=k_descale,
                    v_descale=v_descale,
                    dynamic_causal=dynamic_causal,
                    num_splits=num_splits,
                    s_aux=self.sinks,
                    mask_mod=rswa_mask_mod_fn or mm_mask_mod,
                    aux_tensors=rswa_aux or mm_aux,
                )
                return output

        # Cascade attention (rare case).
        cascade_attention(
            output[:num_actual_tokens],
            query[:num_actual_tokens],
            key_cache,
            value_cache,
            cu_query_lens=attn_metadata.query_start_loc,
            max_query_len=attn_metadata.max_query_len,
            cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,
            prefix_kv_lens=attn_metadata.prefix_kv_lens,
            suffix_kv_lens=attn_metadata.suffix_kv_lens,
            max_kv_len=attn_metadata.max_seq_len,
            softmax_scale=self.scale,
            alibi_slopes=self.alibi_slopes,
            sliding_window=self.sliding_window,
            logits_soft_cap=self.logits_soft_cap,
            block_table=attn_metadata.block_table,
            common_prefix_len=attn_metadata.common_prefix_len,
            max_num_splits=attn_metadata.max_num_splits,
            fa_version=self.vllm_flash_attn_version,
            prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
            suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
            q_descale=layer._q_scale,
            k_descale=layer._k_scale,
            v_descale=layer._v_scale,
            s_aux=self.sinks,
        )
        return output

    def do_kv_cache_update(
        self,
        layer: torch.nn.Module,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> None:
        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
            # For encoder attention,
            # we use direct Q, K, V tensors without caching
            return

        # Scatter write into the KV cache using slot_mapping indices.
        # No TMA kernel is invoked here, so stride canonicalization is not needed.
        # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

        # Reshape the input keys and values and store them in the cache.
        # Skip this if sharing KV cache with an earlier attention layer.
        # NOTE(woosuk): Here, key and value are padded while slot_mapping is
        # not padded. However, we don't need to do key[:num_actual_tokens]
        # and value[:num_actual_tokens] because the reshape_and_cache_flash
        # op uses the slot_mapping's shape to determine the number of
        # actual tokens.
        reshape_and_cache_flash(
            key,
            value,
            key_cache,
            value_cache,
            slot_mapping,
            self.kv_cache_dtype,
            layer._k_scale,
            layer._v_scale,
        )

    def _forward_with_dcp(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        q_descale: torch.Tensor | None = None,
        k_descale: torch.Tensor | None = None,
        v_descale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert self.vllm_flash_attn_version is not None, (
            "FlashAttention version not detected."
        )

        cu_seqlens_q = attn_metadata.query_start_loc
        max_seqlen_q = attn_metadata.max_query_len
        block_table = attn_metadata.block_table

        query = query.contiguous()
        if attn_metadata.max_dcp_context_kv_len == 0:
            flash_attn_varlen_func(
                q=query,
                k=key,
                v=value,
                out=output,
                cu_seqlens_q=cu_seqlens_q,
                max_seqlen_q=max_seqlen_q,
                cu_seqlens_k=cu_seqlens_q,
                max_seqlen_k=max_seqlen_q,
                softmax_scale=self.scale,
                causal=attn_metadata.causal,
                alibi_slopes=self.alibi_slopes,
                window_size=list(self.sliding_window)
                if self.sliding_window is not None
                else None,
                softcap=self.logits_soft_cap,
                return_softmax_lse=True,
                fa_version=self.vllm_flash_attn_version,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
                num_splits=attn_metadata.max_num_splits,
            )
            return output

        query_across_dcp = get_dcp_group().all_gather(query, dim=1)
        sliding_window_size = (
            list(self.sliding_window) if self.sliding_window is not None else None
        )
        n = query_across_dcp.shape[0]
        num_reqs = cu_seqlens_q.shape[0] - 1
        num_decodes = attn_metadata.num_decode_reqs
        num_context_prefills = attn_metadata.num_prefill_reqs
        num_decode_tokens = attn_metadata.num_decode_tokens
        num_context_prefill_tokens = attn_metadata.num_prefill_tokens
        split_dcp_context = should_split_fa2_dcp_context_attention(
            self.vllm_flash_attn_version,
            max_seqlen_q,
            num_reqs,
            num_decodes,
            num_context_prefills,
        )
        dcp_context_out_tokens = max(n, self._dcp_max_num_tokens)
        dcp_context_out_spec = (
            (
                dcp_context_out_tokens,
                self.num_heads * self.dcp_world_size,
                self.head_size,
            ),
            self._dcp_dtype,
        )
        (dcp_context_out_workspace,) = current_workspace_manager().get_simultaneous(
            dcp_context_out_spec,
        )
        dcp_context_out = dcp_context_out_workspace[:n]

        if split_dcp_context:
            # TODO: Remove this DCP + FA2 mixed decode/prefill workaround once
            # FA4 supports this Qwen3.5 shape.
            assert attn_metadata.dcp_context_kv_lens is not None
            assert attn_metadata.max_dcp_context_kv_len is not None
            assert self.vllm_flash_attn_version is not None
            context_attn_out, context_lse = run_split_fa2_dcp_context_attention(
                flash_attn_varlen_func,
                query_across_dcp,
                key_cache,
                value_cache,
                dcp_context_out,
                cu_seqlens_q,
                max_seqlen_q,
                attn_metadata.dcp_context_kv_lens,
                attn_metadata.max_dcp_context_kv_len,
                self.scale,
                self.alibi_slopes,
                sliding_window_size,
                block_table,
                self.logits_soft_cap,
                self.vllm_flash_attn_version,
                q_descale,
                k_descale,
                v_descale,
                attn_metadata.max_num_splits,
                self.num_heads,
                self.dcp_world_size,
                num_decodes,
                num_context_prefills,
                num_decode_tokens,
                num_context_prefill_tokens,
            )
        else:
            context_attn_out, context_lse = flash_attn_varlen_func(
                q=query_across_dcp,
                k=key_cache,
                v=value_cache,
                out=dcp_context_out,
                cu_seqlens_q=cu_seqlens_q,
                max_seqlen_q=max_seqlen_q,
                seqused_k=attn_metadata.dcp_context_kv_lens,
                max_seqlen_k=attn_metadata.max_dcp_context_kv_len,
                softmax_scale=self.scale,
                causal=False,
                alibi_slopes=self.alibi_slopes,
                window_size=sliding_window_size,
                block_table=block_table,
                softcap=self.logits_soft_cap,
                return_softmax_lse=True,
                scheduler_metadata=attn_metadata.scheduler_metadata,
                fa_version=self.vllm_flash_attn_version,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
                num_splits=attn_metadata.max_num_splits,
            )
        # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
        context_attn_out_cor, context_lse_cor = self.dcp_combine(
            context_attn_out,
            context_lse.transpose(0, 1),
            get_dcp_group(),
            return_lse=True,
        )
        context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()

        query_attn_out, query_lse = flash_attn_varlen_func(
            q=query,
            k=key,
            v=value,
            out=output,
            cu_seqlens_q=cu_seqlens_q,
            max_seqlen_q=max_seqlen_q,
            cu_seqlens_k=cu_seqlens_q,
            max_seqlen_k=max_seqlen_q,
            softmax_scale=self.scale,
            causal=attn_metadata.causal,
            alibi_slopes=self.alibi_slopes,
            window_size=sliding_window_size,
            softcap=self.logits_soft_cap,
            return_softmax_lse=True,
            fa_version=self.vllm_flash_attn_version,
            q_descale=q_descale,
            k_descale=k_descale,
            v_descale=v_descale,
            num_splits=attn_metadata.max_num_splits,
        )
        assert context_attn_out_cor.shape == query_attn_out.shape
        assert context_lse_cor.shape == query_lse.shape
        merge_attn_states(
            output,
            context_attn_out_cor,
            context_lse_cor,
            query_attn_out,
            query_lse,
        )

    def _forward_encoder_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        output: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        layer: torch.nn.Module,
    ) -> torch.Tensor:
        """Forward pass for encoder attention without KV cache.

        Args:
            query: shape = [num_encoder_tokens, num_heads, head_size]
            key: shape = [num_encoder_tokens, num_kv_heads, head_size]
            value: shape = [num_encoder_tokens, num_kv_heads, head_size]
            output: shape = [num_encoder_tokens, num_heads, head_size]
            attn_metadata: Encoder attention metadata
            layer: The attention layer

        """
        assert self.vllm_flash_attn_version is not None, (
            "FlashAttention version not detected."
        )

        # For encoder attention, process FP8 quantization if needed
        if is_quantized_kv_cache(self.kv_cache_dtype):
            raise NotImplementedError(
                "quantization is not supported for encoder attention"
            )

        # Use encoder-specific metadata for sequence information
        cu_seqlens_q = attn_metadata.query_start_loc
        cu_seqlens_k = attn_metadata.query_start_loc
        max_seqlen_q = attn_metadata.max_query_len
        max_seqlen_k = attn_metadata.max_query_len

        descale_shape = (
            cu_seqlens_q.shape[0] - 1,  # type: ignore[union-attr]
            self.num_kv_heads,
        )

        # Call flash attention directly on Q, K, V tensors
        sliding_window_size = (
            list(self.sliding_window) if self.sliding_window is not None else None
        )
        flash_attn_varlen_func(
            q=query,
            k=key,
            v=value,
            out=output,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_k=cu_seqlens_k,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_k=max_seqlen_k,
            softmax_scale=self.scale,
            causal=False,  # Encoder attention is bidirectional
            alibi_slopes=self.alibi_slopes,
            window_size=sliding_window_size,
            softcap=self.logits_soft_cap,
            fa_version=self.vllm_flash_attn_version,
            q_descale=layer._q_scale.expand(descale_shape)  # type: ignore[operator]
            if self.supports_quant_query_input
            else None,
            k_descale=layer._k_scale.expand(descale_shape),  # type: ignore[operator]
            v_descale=layer._v_scale.expand(descale_shape),  # type: ignore[operator]
            # The hd256 kernel does not support SplitKV.
            num_splits=1 if self.batch_invariant_enabled or self.fa4_hd256 else 0,
            s_aux=self.sinks,
        )

        return output

_forward_encoder_attention(query, key, value, output, attn_metadata, layer)

Forward pass for encoder attention without KV cache.

Parameters:

• query

  (Tensor) –

  shape = [num_encoder_tokens, num_heads, head_size]

• key

  (Tensor) –

  shape = [num_encoder_tokens, num_kv_heads, head_size]

• value

  (Tensor) –

  shape = [num_encoder_tokens, num_kv_heads, head_size]

• output

  (Tensor) –

  shape = [num_encoder_tokens, num_heads, head_size]

• attn_metadata

  (FlashAttentionMetadata) –

  Encoder attention metadata

• layer

  (Module) –

  The attention layer

Source code in vllm/v1/attention/backends/flash_attn.py

def _forward_encoder_attention(
    self,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    attn_metadata: FlashAttentionMetadata,
    layer: torch.nn.Module,
) -> torch.Tensor:
    """Forward pass for encoder attention without KV cache.

    Args:
        query: shape = [num_encoder_tokens, num_heads, head_size]
        key: shape = [num_encoder_tokens, num_kv_heads, head_size]
        value: shape = [num_encoder_tokens, num_kv_heads, head_size]
        output: shape = [num_encoder_tokens, num_heads, head_size]
        attn_metadata: Encoder attention metadata
        layer: The attention layer

    """
    assert self.vllm_flash_attn_version is not None, (
        "FlashAttention version not detected."
    )

    # For encoder attention, process FP8 quantization if needed
    if is_quantized_kv_cache(self.kv_cache_dtype):
        raise NotImplementedError(
            "quantization is not supported for encoder attention"
        )

    # Use encoder-specific metadata for sequence information
    cu_seqlens_q = attn_metadata.query_start_loc
    cu_seqlens_k = attn_metadata.query_start_loc
    max_seqlen_q = attn_metadata.max_query_len
    max_seqlen_k = attn_metadata.max_query_len

    descale_shape = (
        cu_seqlens_q.shape[0] - 1,  # type: ignore[union-attr]
        self.num_kv_heads,
    )

    # Call flash attention directly on Q, K, V tensors
    sliding_window_size = (
        list(self.sliding_window) if self.sliding_window is not None else None
    )
    flash_attn_varlen_func(
        q=query,
        k=key,
        v=value,
        out=output,
        cu_seqlens_q=cu_seqlens_q,
        cu_seqlens_k=cu_seqlens_k,
        max_seqlen_q=max_seqlen_q,
        max_seqlen_k=max_seqlen_k,
        softmax_scale=self.scale,
        causal=False,  # Encoder attention is bidirectional
        alibi_slopes=self.alibi_slopes,
        window_size=sliding_window_size,
        softcap=self.logits_soft_cap,
        fa_version=self.vllm_flash_attn_version,
        q_descale=layer._q_scale.expand(descale_shape)  # type: ignore[operator]
        if self.supports_quant_query_input
        else None,
        k_descale=layer._k_scale.expand(descale_shape),  # type: ignore[operator]
        v_descale=layer._v_scale.expand(descale_shape),  # type: ignore[operator]
        # The hd256 kernel does not support SplitKV.
        num_splits=1 if self.batch_invariant_enabled or self.fa4_hd256 else 0,
        s_aux=self.sinks,
    )

    return output

forward(layer, query, key, value, kv_cache, attn_metadata, output, output_scale=None, output_block_scale=None)

Forward pass with FlashAttention.

Parameters:

• layer

  (Module) –

  The attention layer, providing the q/k/v quantization scales.

• query

  (Tensor) –

  shape = [num_tokens, num_heads, head_size]

• key

  (Tensor) –

  shape = [num_tokens, num_kv_heads, head_size]

• value

  (Tensor) –

  shape = [num_tokens, num_kv_heads, head_size]

• kv_cache

  (Tensor) –

  shape = [num_blocks, num_kv_heads, block_size, 2 * head_size]

• attn_metadata

  (FlashAttentionMetadata) –

  Metadata for attention.

• output

  (Tensor) –

  Tensor that the attention result is written into.

• output_scale

  (Tensor | None, default: None ) –

  Scale for fused output quantization; not supported by this backend.

• output_block_scale

  (Tensor | None, default: None ) –

  Block scale for fused output quantization; not supported by this backend.

Returns:

• Tensor –

  shape = [num_tokens, num_heads * head_size]

NOTE: FP8 quantization, flash-attn expect the size of {q,k,v}_descale to be (num_sequences, num_kv_heads). We use torch's .expand() to avoid duplicating values

Source code in vllm/v1/attention/backends/flash_attn.py

def forward(
    self,
    layer: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: FlashAttentionMetadata,
    output: torch.Tensor,
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
) -> torch.Tensor:
    """Forward pass with FlashAttention.

    Args:
        layer: The attention layer, providing the q/k/v quantization scales.
        query: shape = [num_tokens, num_heads, head_size]
        key: shape = [num_tokens, num_kv_heads, head_size]
        value: shape = [num_tokens, num_kv_heads, head_size]
        kv_cache: shape =
            [num_blocks, num_kv_heads, block_size, 2 * head_size]
        attn_metadata: Metadata for attention.
        output: Tensor that the attention result is written into.
        output_scale: Scale for fused output quantization; not supported
            by this backend.
        output_block_scale: Block scale for fused output quantization;
            not supported by this backend.

    Returns:
        shape = [num_tokens, num_heads * head_size]
    NOTE: FP8 quantization, flash-attn expect the size of
          {q,k,v}_descale to be (num_sequences, num_kv_heads).
          We use torch's .expand() to avoid duplicating values

    """
    assert self.vllm_flash_attn_version is not None, (
        "FlashAttention version not detected."
    )

    if output_scale is not None or output_block_scale is not None:
        raise NotImplementedError(
            "fused output quantization is not yet supported for FlashAttentionImpl"
        )

    if attn_metadata is None:
        # Profiling run.
        return output.fill_(0)

    attn_type = self.attn_type

    # IMPORTANT!
    # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
    # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
    # in this method. For example, `view` and `slice` (or `[:n]`) operations
    # are surprisingly slow even in the case they do not invoke any GPU ops.
    # Minimize the PyTorch ops in this method as much as possible.
    # Whenever making a change in this method, please benchmark the
    # performance to make sure it does not introduce any overhead.

    num_actual_tokens = attn_metadata.num_actual_tokens

    # Handle encoder attention differently - no KV cache needed
    if attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
        # For encoder attention,
        # we use direct Q, K, V tensors without caching
        return self._forward_encoder_attention(
            query[:num_actual_tokens],
            key[:num_actual_tokens],
            value[:num_actual_tokens],
            output[:num_actual_tokens],
            attn_metadata,
            layer,
        )

    # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
    key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
    # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
    # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
    # See vllm.utils.torch_utils.canonicalize_singleton_dim_strides.
    fixed_k = canonicalize_singleton_dim_strides(key_cache)
    fixed_v = canonicalize_singleton_dim_strides(value_cache)
    if fixed_k is not key_cache or fixed_v is not value_cache:
        logger.debug(
            "Canonicalized degenerate KV cache strides (FlashAttention): "
            "shape=%s, key strides before=%s after=%s, "
            "value strides before=%s after=%s",
            key_cache.shape,
            key_cache.stride(),
            fixed_k.stride(),
            value_cache.stride(),
            fixed_v.stride(),
        )
    key_cache, value_cache = fixed_k, fixed_v

    if is_quantized_kv_cache(self.kv_cache_dtype):
        # queries are quantized in the attention layer
        key_cache = key_cache.view(current_platform.fp8_dtype())
        value_cache = value_cache.view(current_platform.fp8_dtype())

    if not attn_metadata.use_cascade:
        cu_seqlens_q = attn_metadata.query_start_loc
        seqused_k = attn_metadata.seq_lens
        max_seqlen_q = attn_metadata.max_query_len
        max_seqlen_k = attn_metadata.max_seq_len
        block_table = attn_metadata.block_table
        scheduler_metadata = attn_metadata.scheduler_metadata

        descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)

        q_descale = (
            layer._q_scale.expand(descale_shape)
            if self.supports_quant_query_input
            else None
        )
        k_descale = layer._k_scale.expand(descale_shape)
        v_descale = layer._v_scale.expand(descale_shape)

        if self.dcp_world_size > 1:
            self._forward_with_dcp(
                query[:num_actual_tokens],
                key[:num_actual_tokens],
                value[:num_actual_tokens],
                key_cache,
                value_cache,
                output[:num_actual_tokens],
                attn_metadata,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
            )
            return output
        else:
            causal = attn_metadata.causal
            is_dynamic_causal = isinstance(causal, torch.Tensor)

            # The layer's own window wins over the group's: one KV cache
            # group can hold both windowed and global layers (e.g. Gemma-3
            # with the hybrid KV cache manager disabled), and the group spec
            # cannot describe both.
            window = _maybe_symmetrize_window(self.sliding_window, causal)
            sliding_window_size: list[int] | None = (
                list(window) if window is not None else None
            )

            mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
            mm_mask_mod = None
            mm_aux = None
            if (
                mm_prefix_query_ranges is not None
                and not is_dynamic_causal
                and causal is True
                and self.vllm_flash_attn_version == 4
            ):
                # Triton convention: 1 + window_size[0]. Global layers store
                # (-1, -1) → sw stays None.
                layer_window = self.sliding_window
                sw_val = (
                    1 + layer_window[0]
                    if layer_window is not None and layer_window[0] >= 0
                    else None
                )
                # Gemma4: also clamp the bidirectional block to the
                # sliding window when the layer opts in
                # (mm_prefix_clamp_sliding_window flag from PR #47217).
                mm_clamp_sw = 0
                if (
                    getattr(layer, "mm_prefix_clamp_sliding_window", False)
                    and sw_val is not None
                ):
                    mm_clamp_sw = sw_val
                mm_mask_mod = _make_mm_prefix_mask_mod(
                    sliding_window=mm_clamp_sw,
                    sliding_window_left=sw_val,
                )
                mm_aux = [mm_prefix_query_ranges, attn_metadata.query_start_loc]
                # mm_prefix is (causal ∧ window) ∨ bidirectional-range —
                # not ⊆ causal. FA #155 stopped auto-clearing causal/local
                # when mask_mod is set, so the caller must disable them or
                # the built-in causal path shorts out / clips the mask_mod.
                causal = False
                sliding_window_size = None

            # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level
            # mask without block-size approximation.  The mask_mod encodes
            # "causal AND (kv < prefix_len OR q - kv < rswa_window)", which
            # supersedes any FA-layer sliding_window_size parameter.
            rswa_mask_mod_fn = None
            rswa_aux = None
            if (
                attn_metadata.rswa_prefix_lens is not None
                and self.vllm_flash_attn_version == 4
                and not is_dynamic_causal
            ):
                rswa_mask_mod_fn = _make_rswa_mask_mod()
                rswa_aux = [
                    attn_metadata.rswa_prefix_lens.to(torch.int32),
                    attn_metadata.rswa_window_tensor,  # pre-allocated CUDA tensor
                ]
                # mask_mod fully expresses R-SWA; disable FA's own window.
                sliding_window_size = None

            dynamic_causal = None
            if isinstance(causal, torch.Tensor):
                if self.vllm_flash_attn_version != 4:
                    raise NotImplementedError(
                        "Per-sequence causal requires FA4. Current version: "
                        f"FA{self.vllm_flash_attn_version}"
                    )
                dynamic_causal = causal
                has_window = (
                    sliding_window_size is not None and sliding_window_size[1] >= 0
                )
                causal = not has_window

            num_splits = attn_metadata.max_num_splits
            if self.fa4_hd256:
                # hd256 requires page-aligned lengths, exact-width block
                # tables, and no SplitKV.
                num_pages = cdiv(max_seqlen_k, FA4_HD256_PAGE_SIZE)
                max_seqlen_k = num_pages * FA4_HD256_PAGE_SIZE
                block_table = block_table[:, :num_pages]
                num_splits = 1

            _FA4_DENSE_ATTENTION_KERNEL(
                q=query[:num_actual_tokens],
                k=key_cache,
                v=value_cache,
                out=output[:num_actual_tokens],
                cu_seqlens_q=cu_seqlens_q,
                max_seqlen_q=max_seqlen_q,
                seqused_k=seqused_k,
                max_seqlen_k=max_seqlen_k,
                softmax_scale=self.scale,
                causal=causal,
                alibi_slopes=self.alibi_slopes,
                window_size=sliding_window_size,
                block_table=block_table,
                softcap=self.logits_soft_cap,
                scheduler_metadata=scheduler_metadata,
                fa_version=self.vllm_flash_attn_version,
                q_descale=q_descale,
                k_descale=k_descale,
                v_descale=v_descale,
                dynamic_causal=dynamic_causal,
                num_splits=num_splits,
                s_aux=self.sinks,
                mask_mod=rswa_mask_mod_fn or mm_mask_mod,
                aux_tensors=rswa_aux or mm_aux,
            )
            return output

    # Cascade attention (rare case).
    cascade_attention(
        output[:num_actual_tokens],
        query[:num_actual_tokens],
        key_cache,
        value_cache,
        cu_query_lens=attn_metadata.query_start_loc,
        max_query_len=attn_metadata.max_query_len,
        cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,
        prefix_kv_lens=attn_metadata.prefix_kv_lens,
        suffix_kv_lens=attn_metadata.suffix_kv_lens,
        max_kv_len=attn_metadata.max_seq_len,
        softmax_scale=self.scale,
        alibi_slopes=self.alibi_slopes,
        sliding_window=self.sliding_window,
        logits_soft_cap=self.logits_soft_cap,
        block_table=attn_metadata.block_table,
        common_prefix_len=attn_metadata.common_prefix_len,
        max_num_splits=attn_metadata.max_num_splits,
        fa_version=self.vllm_flash_attn_version,
        prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
        suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
        q_descale=layer._q_scale,
        k_descale=layer._k_scale,
        v_descale=layer._v_scale,
        s_aux=self.sinks,
    )
    return output

FlashAttentionMetadataBuilder

Bases: AttentionMetadataBuilder[FlashAttentionMetadata]

Methods:

• build –

  fast_build disables AOT scheduling, used when there will be few

Source code in vllm/v1/attention/backends/flash_attn.py

class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetadata]):
    # FA3:
    # Supports full cudagraphs for all cases.
    #
    # FA2:
    # For FA2, a graph is captured with max_query_len=1, (which is what we
    # capture by default for num_tokens <= max_num_seqs when there is no
    # spec-decode) then these graphs will not work for mixed prefill-decode
    # (unlike FA3). This is due to special max_query_len=1 packed-GQA handling
    # in FA2.
    # In summary if we are running with spec decodes the graphs would
    # work for mixed prefill-decode and uniform-decode. But for non-spec decodes
    # the graphs would not work for mixed prefill-decode; sorta the inverse
    # of UNIFORM_SINGLE_TOKEN_DECODE.
    # There's probably a better way to describe this using `AttentionCGSupport`
    # but for now just set it to `UNIFORM_BATCH` to get use to drop down
    # to FULL_AND_PIECEWISE.
    # TODO(luka, lucas): audit FA2 as part of:
    #  https://github.com/vllm-project/vllm/issues/22945
    _cudagraph_support = (
        AttentionCGSupport.ALWAYS
        if get_flash_attn_version() == 3
        else AttentionCGSupport.UNIFORM_BATCH
    )
    supports_update_block_table: bool = True

    @classmethod
    def get_cudagraph_support(
        cls,
        vllm_config: "VllmConfig",
        kv_cache_spec: "KVCacheSpec",
    ) -> AttentionCGSupport:
        return cls._cudagraph_support

    def _get_scheduler_metadata(
        self,
        *,
        aot_schedule: bool,
        batch_size: int,
        cu_query_lens: torch.Tensor,
        max_query_len: int,
        seqlens: torch.Tensor,
        max_seq_len: int,
        causal: bool | torch.Tensor,
        max_num_splits: int,
    ) -> torch.Tensor | None:
        if not aot_schedule:
            return None

        cache_dtype = self.cache_config.cache_dtype
        if is_quantized_kv_cache(cache_dtype):
            qkv_dtype = current_platform.fp8_dtype()
        else:
            qkv_dtype = self.kv_cache_dtype
        return get_scheduler_metadata(
            batch_size=batch_size,
            max_seqlen_q=max_query_len,
            max_seqlen_k=max_seq_len,
            num_heads_q=self.num_heads_q * self.dcp_world_size,
            num_heads_kv=self.num_heads_kv,
            headdim=self.headdim,
            cache_seqlens=seqlens,
            qkv_dtype=qkv_dtype,
            cu_seqlens_q=cu_query_lens,
            page_size=self.block_size,
            causal=causal,
            window_size=_maybe_symmetrize_window(self.aot_sliding_window, causal),
            num_splits=max_num_splits,
        )

    def _store_scheduler_metadata(
        self, scheduler_metadata: torch.Tensor | None
    ) -> torch.Tensor | None:
        if self.use_full_cuda_graph and scheduler_metadata is not None:
            n = scheduler_metadata.shape[0]
            assert self.scheduler_metadata is not None
            self.scheduler_metadata[:n] = scheduler_metadata
            # NOTE(woosuk): We should zero out the rest of the scheduler
            # metadata to guarantee the correctness. Otherwise, some thread
            # blocks may use the invalid scheduler metadata and overwrite the
            # output buffer.
            self.scheduler_metadata[n:] = 0
            return self.scheduler_metadata[:n]
        return scheduler_metadata

    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
        self.model_config = vllm_config.model_config
        self.parallel_config = vllm_config.parallel_config
        self.cache_config = vllm_config.cache_config
        self.compilation_config = vllm_config.compilation_config
        self.attention_config = vllm_config.attention_config

        self.num_heads_q = get_num_attention_heads_from_layers(
            vllm_config, layer_names
        ) or self.model_config.get_num_attention_heads(self.parallel_config)
        self.num_heads_kv = kv_cache_spec.num_kv_heads
        self.kv_cache_dtype = kv_cache_spec.dtype
        self.headdim = kv_cache_spec.head_size
        self.block_size = kv_cache_spec.block_size

        self.max_num_splits = 0  # No upper bound on the number of splits.
        self.aot_schedule = get_flash_attn_version() == 3
        head_size_v = getattr(kv_cache_spec, "head_size_v", None)
        fa_version = get_flash_attn_version(
            head_size=self.headdim,
            head_size_v=head_size_v,
            kv_cache_block_size=self.block_size,
            supports_fa4_hd256=True,
        )
        self.fa4_hd256 = fa_version == 4 and uses_fa4_hd256_kernel(
            self.headdim, head_size_v
        )

        if (
            vllm_config.kernel_config.enable_jit_warmup
            and self.model_config.hf_config.model_type in ("gemma4", "gemma4_unified")
            and fa_version == 4
        ):
            _FA4_DENSE_ATTENTION_KERNEL.register_warmup(
                vllm_config=vllm_config,
                kv_cache_spec=kv_cache_spec,
                num_qo_heads=self.num_heads_q,
            )

        self.dcp_world_size, self.dcp_rank = get_dcp_world_size_and_rank(
            kv_cache_spec.dcp_sharded
        )

        # Fused draft decode reuses the captured metadata object across draft
        # steps. For DCP, build-time host-side decisions such as
        # skip_dcp_context_attention() can change the metadata shape/control
        # path (for example max_dcp_context_kv_len), and those Python-side
        # fields are not refreshed in-place between graph replays. Keep the
        # fused path disabled until DCP gets a full replay-safe refresh model.
        self.supports_draft_decode_metadata_update = self.dcp_world_size == 1

        self.cp_kv_cache_interleave_size = (
            self.parallel_config.cp_kv_cache_interleave_size
        )

        self.use_full_cuda_graph = (
            self.compilation_config.cudagraph_mode.has_full_cudagraphs()
        )
        self.max_cudagraph_size = self.compilation_config.max_cudagraph_capture_size

        if self.use_full_cuda_graph and self.aot_schedule:
            # FA3 scheduler_metadata size: 1 + round_up(batch_size, 4) * 4
            # The +1 is for the tile_count_semaphore (synchronization).
            # The 4 slots per batch element (num_prepare_batch_vectors) are:
            #   prepare_varlen + dynamic_split + sort_batches + head_swizzle
            # See: https://github.com/vllm-project/flash-attention/blob/5824e6e/hopper/flash_api.cpp#L664-L671  # noqa: E501
            max_batch_size = max(
                vllm_config.scheduler_config.max_num_seqs,
                self.max_cudagraph_size or 0,
            )
            self.scheduler_metadata = torch.zeros(
                1 + round_up(max_batch_size, 4) * 4,
                dtype=torch.int32,
                device=self.device,
            )
            # When using cuda graph, we need to set the upper bound of the
            # number of splits so that large enough intermediate buffers are
            # pre-allocated during capture.
            self.max_num_splits = (
                self.attention_config.flash_attn_max_num_splits_for_cuda_graph
            )

        if self.dcp_world_size > 1:
            max_num_reqs = vllm_config.scheduler_config.max_num_seqs
            self._dcp_context_kv_lens = torch.zeros(
                max_num_reqs,
                dtype=torch.int32,
                device=self.device,
            )

        # Sliding window size to be used with the AOT scheduler will be
        # populated on first build() call.
        self.aot_sliding_window: tuple[int, int] | None = None

        # R-SWA: persistent CUDA-graph-safe buffers owned by this builder.
        self.rswa_window: int | None = self.model_config.rswa_window
        self.persistent_rswa_prefix_lens: torch.Tensor | None = None
        self.persistent_rswa_window_tensor: torch.Tensor | None = None
        if self.rswa_window is not None:
            max_num_reqs = vllm_config.scheduler_config.max_num_seqs
            self.persistent_rswa_prefix_lens = torch.zeros(
                max_num_reqs, dtype=torch.int32, device=self.device
            )
            self.persistent_rswa_window_tensor = torch.tensor(
                [self.rswa_window], dtype=torch.int32, device=self.device
            )

        # mm_prefix: persistent staging + device buffers owned by this builder,
        # sized by scheduled query tokens so build() never allocates.
        self.mm_prefix_query_ranges_cpu: torch.Tensor | None = None
        self.mm_prefix_query_ranges_np: np.ndarray | None = None
        self.mm_prefix_query_ranges_gpu: torch.Tensor | None = None
        if self.model_config.is_mm_prefix_lm:
            max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
            self.mm_prefix_query_ranges_cpu = torch.empty(
                (max_num_tokens, 2), dtype=torch.int32, pin_memory=PIN_MEMORY
            )
            self.mm_prefix_query_ranges_np = self.mm_prefix_query_ranges_cpu.numpy()
            self.mm_prefix_query_ranges_gpu = torch.empty(
                (max_num_tokens, 2), dtype=torch.int32, device=self.device
            )

    def build(
        self,
        common_prefix_len: int,
        common_attn_metadata: CommonAttentionMetadata,
        fast_build: bool = False,
    ) -> FlashAttentionMetadata:
        """fast_build disables AOT scheduling, used when there will be few
        iterations i.e. spec-decode
        """
        num_reqs = common_attn_metadata.num_reqs
        num_actual_tokens = common_attn_metadata.num_actual_tokens
        max_query_len = common_attn_metadata.max_query_len
        max_seq_len = common_attn_metadata.max_seq_len
        query_start_loc = common_attn_metadata.query_start_loc
        seq_lens = common_attn_metadata.seq_lens
        block_table_tensor = common_attn_metadata.block_table_tensor
        slot_mapping = common_attn_metadata.slot_mapping
        causal = common_attn_metadata.causal

        # Disable AOT schedule for spec-decode proposer (not worth the overhead)
        # and for batch invariance (schedule varies with max_seqlen_q/k).
        aot_schedule = (
            self.aot_schedule and not fast_build and not envs.VLLM_BATCH_INVARIANT
        )

        if self.aot_sliding_window is None:
            self.aot_sliding_window = (-1, -1)
            # For the AOT scheduler we need the sliding window value to be
            # constant for all layers to. We have to populate this on the first
            # build() call so the layers are constructed (cannot populate)
            # in __init__.
            if aot_schedule:
                sliding_window_configs = _get_sliding_window_configs(self.vllm_config)
                if len(sliding_window_configs) == 1:
                    sliding_window_config = sliding_window_configs.pop()
                    if sliding_window_config is not None:
                        self.aot_sliding_window = sliding_window_config
                elif len(sliding_window_configs) > 1:
                    self.aot_schedule = False
                    aot_schedule = False

        max_num_splits = 0  # 0 means use FA3's heuristics, not CG compatible
        if (
            self.use_full_cuda_graph
            and self.max_cudagraph_size is not None
            and num_actual_tokens <= self.max_cudagraph_size
        ):
            # NOTE(woosuk): Setting num_splits > 1 may increase the memory
            # usage, because the intermediate buffers of size [num_splits,
            # num_heads, num_tokens, head_size] are allocated. Therefore,
            # we only set num_splits when using cuda graphs.
            max_num_splits = self.max_num_splits

        if envs.VLLM_BATCH_INVARIANT:
            max_num_splits = 1

        use_cascade = common_prefix_len > 0
        max_dcp_context_kv_len = 0
        dcp_context_kv_lens = None
        num_decode_reqs = 0
        num_prefill_reqs = 0
        num_decode_tokens = 0
        num_prefill_tokens = 0

        cu_prefix_query_lens = None
        prefix_kv_lens = None
        suffix_kv_lens = None
        prefix_scheduler_metadata = None

        if self.dcp_world_size > 1:
            query_lens = query_start_loc[1:] - query_start_loc[:-1]
            context_kv_lens = seq_lens - query_lens
            local_context_kv_lens = get_dcp_local_seq_lens(
                context_kv_lens,
                self.dcp_world_size,
                self.dcp_rank,
                self.cp_kv_cache_interleave_size,
            )
            self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens
            self._dcp_context_kv_lens[num_reqs:] = 0
            dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs]

            skip_dcp_context_attention = False
            if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
                query_lens_cpu = (
                    common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1]
                    - common_attn_metadata.query_start_loc_cpu[:num_reqs]
                )
                context_kv_lens_cpu = (
                    common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs]
                    - query_lens_cpu
                )
                skip_dcp_context_attention = should_skip_dcp_context_attention(
                    context_kv_lens_cpu
                )

            if max_query_len > 1:
                (
                    num_decode_reqs,
                    num_prefill_reqs,
                    num_decode_tokens,
                    num_prefill_tokens,
                ) = split_dcp_context_queries(
                    common_attn_metadata.query_start_loc_cpu,
                    common_attn_metadata.seq_lens_cpu_upper_bound,
                    max_query_len,
                    num_actual_tokens,
                )

            # After DCP distribution, the maximum number of tokens for any rank is
            # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
            # and I is cp_kv_cache_interleave_size.
            # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
            if skip_dcp_context_attention:
                max_dcp_context_kv_len = 0
                scheduler_metadata = None
            else:
                num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
                max_dcp_context_kv_len = (
                    (max_seq_len + num_partitions - 1) // num_partitions
                ) * self.cp_kv_cache_interleave_size

                scheduler_metadata = self._get_scheduler_metadata(
                    aot_schedule=aot_schedule,
                    batch_size=num_reqs,
                    cu_query_lens=query_start_loc,
                    max_query_len=max_query_len,
                    seqlens=dcp_context_kv_lens,
                    max_seq_len=max_dcp_context_kv_len,
                    causal=False,
                    max_num_splits=max_num_splits,
                )
        elif use_cascade:
            cu_prefix_query_lens = torch.tensor(
                [0, num_actual_tokens], dtype=torch.int32, device=self.device
            )
            prefix_kv_lens = torch.tensor(
                [common_prefix_len], dtype=torch.int32, device=self.device
            )
            # Use GPU tensor directly - no CPU sync needed
            suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len
            prefix_scheduler_metadata = self._get_scheduler_metadata(
                aot_schedule=aot_schedule,
                batch_size=1,
                cu_query_lens=cu_prefix_query_lens,
                max_query_len=num_actual_tokens,
                seqlens=prefix_kv_lens,
                max_seq_len=common_prefix_len,
                causal=False,
                max_num_splits=max_num_splits,
            )
            scheduler_metadata = self._get_scheduler_metadata(
                aot_schedule=aot_schedule,
                batch_size=num_reqs,
                cu_query_lens=query_start_loc,
                max_query_len=max_query_len,
                seqlens=suffix_kv_lens,
                max_seq_len=max_seq_len - common_prefix_len,
                causal=True,
                max_num_splits=max_num_splits,
            )
        else:
            scheduler_metadata = self._get_scheduler_metadata(
                aot_schedule=aot_schedule,
                batch_size=num_reqs,
                cu_query_lens=query_start_loc,
                max_query_len=max_query_len,
                seqlens=seq_lens,
                max_seq_len=max_seq_len,
                causal=causal,
                max_num_splits=max_num_splits,
            )
        scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)

        if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
            raise ValueError(
                f"Per-request causal tensor must be int32, got {causal.dtype}. "
                "Casting here would allocate a fresh tensor each build() and "
                "break FULL CUDA graph replay; allocate the buffer as int32."
            )

        attn_metadata = FlashAttentionMetadata(
            num_actual_tokens=num_actual_tokens,
            max_query_len=max_query_len,
            query_start_loc=query_start_loc,
            max_seq_len=max_seq_len,
            seq_lens=seq_lens,
            block_table=block_table_tensor,
            slot_mapping=slot_mapping,
            max_dcp_context_kv_len=max_dcp_context_kv_len,
            dcp_context_kv_lens=dcp_context_kv_lens,
            num_decode_reqs=num_decode_reqs,
            num_prefill_reqs=num_prefill_reqs,
            num_decode_tokens=num_decode_tokens,
            num_prefill_tokens=num_prefill_tokens,
            use_cascade=use_cascade,
            common_prefix_len=common_prefix_len,
            scheduler_metadata=scheduler_metadata,
            cu_prefix_query_lens=cu_prefix_query_lens,
            prefix_kv_lens=prefix_kv_lens,
            suffix_kv_lens=suffix_kv_lens,
            prefix_scheduler_metadata=prefix_scheduler_metadata,
            max_num_splits=max_num_splits,
            causal=causal,
        )

        # Compute mm_prefix range tensor if the batch contains
        # multimodal tokens with bidirectional ranges.  Built for every FA
        # group; Gemma4 nulls the field for its non-sliding layers.
        mm_ranges = common_attn_metadata.mm_req_doc_ranges
        if mm_ranges is not None and self.mm_prefix_query_ranges_np is not None:
            # The upper bound is exact for prefill rows, which is where
            # mm_prefix ranges live; decode rows only ever get an optimistic
            # (larger) context, moving them further past every range.
            assert common_attn_metadata.seq_lens_cpu_upper_bound is not None, (
                "mm_prefix requires seq_lens_cpu_upper_bound"
            )
            num_mm_tokens = fill_mm_prefix_query_ranges(
                self.mm_prefix_query_ranges_np,
                mm_ranges,
                common_attn_metadata.query_start_loc_cpu,
                common_attn_metadata.seq_lens_cpu_upper_bound,
            )
            if num_mm_tokens > 0:
                assert self.mm_prefix_query_ranges_cpu is not None
                assert self.mm_prefix_query_ranges_gpu is not None
                mm_query_ranges = self.mm_prefix_query_ranges_gpu[:num_mm_tokens]
                mm_query_ranges.copy_(
                    self.mm_prefix_query_ranges_cpu[:num_mm_tokens],
                    non_blocking=True,
                )
                attn_metadata.mm_prefix_query_range_tensor = mm_query_ranges

        # R-SWA: copy prefix lengths into persistent buffers (outside the
        # compiled region) so forward() never allocates during CUDA graph
        # capture.  rswa_window is a static model config scalar read here.
        if (
            self.rswa_window is not None
            and common_attn_metadata.rswa_prefix_lens is not None
        ):
            assert self.persistent_rswa_prefix_lens is not None
            assert self.persistent_rswa_window_tensor is not None
            src = common_attn_metadata.rswa_prefix_lens
            rswa_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]
            rswa_prefix_lens.copy_(src[:num_reqs], non_blocking=True)
            attn_metadata.rswa_prefix_lens = rswa_prefix_lens
            attn_metadata.rswa_window = self.rswa_window
            attn_metadata.rswa_window_tensor = self.persistent_rswa_window_tensor

        return attn_metadata

    def update_block_table(
        self,
        metadata: FlashAttentionMetadata,
        blk_table: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> FlashAttentionMetadata:
        new_metadata = copy.copy(metadata)
        new_metadata.block_table = blk_table
        new_metadata.slot_mapping = slot_mapping
        return new_metadata

    def update_draft_decode_metadata(self, metadata: FlashAttentionMetadata) -> None:
        if metadata.scheduler_metadata is None:
            return

        num_reqs = metadata.num_decode_reqs or metadata.seq_lens.shape[0]

        assert self.dcp_world_size == 1
        assert not metadata.use_cascade

        scheduler_metadata = self._get_scheduler_metadata(
            aot_schedule=True,
            batch_size=num_reqs,
            cu_query_lens=metadata.query_start_loc,
            max_query_len=metadata.max_query_len,
            seqlens=metadata.seq_lens,
            max_seq_len=metadata.max_seq_len,
            causal=metadata.causal,
            max_num_splits=metadata.max_num_splits,
        )

        metadata.scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)

    def use_cascade_attention(self, *args, **kwargs) -> bool:
        if self.fa4_hd256:
            # Cascade may use a non-page-aligned prefix length.
            return False
        return use_cascade_attention(*args, **kwargs)

build(common_prefix_len, common_attn_metadata, fast_build=False)

fast_build disables AOT scheduling, used when there will be few iterations i.e. spec-decode

Source code in vllm/v1/attention/backends/flash_attn.py

def build(
    self,
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> FlashAttentionMetadata:
    """fast_build disables AOT scheduling, used when there will be few
    iterations i.e. spec-decode
    """
    num_reqs = common_attn_metadata.num_reqs
    num_actual_tokens = common_attn_metadata.num_actual_tokens
    max_query_len = common_attn_metadata.max_query_len
    max_seq_len = common_attn_metadata.max_seq_len
    query_start_loc = common_attn_metadata.query_start_loc
    seq_lens = common_attn_metadata.seq_lens
    block_table_tensor = common_attn_metadata.block_table_tensor
    slot_mapping = common_attn_metadata.slot_mapping
    causal = common_attn_metadata.causal

    # Disable AOT schedule for spec-decode proposer (not worth the overhead)
    # and for batch invariance (schedule varies with max_seqlen_q/k).
    aot_schedule = (
        self.aot_schedule and not fast_build and not envs.VLLM_BATCH_INVARIANT
    )

    if self.aot_sliding_window is None:
        self.aot_sliding_window = (-1, -1)
        # For the AOT scheduler we need the sliding window value to be
        # constant for all layers to. We have to populate this on the first
        # build() call so the layers are constructed (cannot populate)
        # in __init__.
        if aot_schedule:
            sliding_window_configs = _get_sliding_window_configs(self.vllm_config)
            if len(sliding_window_configs) == 1:
                sliding_window_config = sliding_window_configs.pop()
                if sliding_window_config is not None:
                    self.aot_sliding_window = sliding_window_config
            elif len(sliding_window_configs) > 1:
                self.aot_schedule = False
                aot_schedule = False

    max_num_splits = 0  # 0 means use FA3's heuristics, not CG compatible
    if (
        self.use_full_cuda_graph
        and self.max_cudagraph_size is not None
        and num_actual_tokens <= self.max_cudagraph_size
    ):
        # NOTE(woosuk): Setting num_splits > 1 may increase the memory
        # usage, because the intermediate buffers of size [num_splits,
        # num_heads, num_tokens, head_size] are allocated. Therefore,
        # we only set num_splits when using cuda graphs.
        max_num_splits = self.max_num_splits

    if envs.VLLM_BATCH_INVARIANT:
        max_num_splits = 1

    use_cascade = common_prefix_len > 0
    max_dcp_context_kv_len = 0
    dcp_context_kv_lens = None
    num_decode_reqs = 0
    num_prefill_reqs = 0
    num_decode_tokens = 0
    num_prefill_tokens = 0

    cu_prefix_query_lens = None
    prefix_kv_lens = None
    suffix_kv_lens = None
    prefix_scheduler_metadata = None

    if self.dcp_world_size > 1:
        query_lens = query_start_loc[1:] - query_start_loc[:-1]
        context_kv_lens = seq_lens - query_lens
        local_context_kv_lens = get_dcp_local_seq_lens(
            context_kv_lens,
            self.dcp_world_size,
            self.dcp_rank,
            self.cp_kv_cache_interleave_size,
        )
        self._dcp_context_kv_lens[:num_reqs] = local_context_kv_lens
        self._dcp_context_kv_lens[num_reqs:] = 0
        dcp_context_kv_lens = self._dcp_context_kv_lens[:num_reqs]

        skip_dcp_context_attention = False
        if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
            query_lens_cpu = (
                common_attn_metadata.query_start_loc_cpu[1 : num_reqs + 1]
                - common_attn_metadata.query_start_loc_cpu[:num_reqs]
            )
            context_kv_lens_cpu = (
                common_attn_metadata.seq_lens_cpu_upper_bound[:num_reqs]
                - query_lens_cpu
            )
            skip_dcp_context_attention = should_skip_dcp_context_attention(
                context_kv_lens_cpu
            )

        if max_query_len > 1:
            (
                num_decode_reqs,
                num_prefill_reqs,
                num_decode_tokens,
                num_prefill_tokens,
            ) = split_dcp_context_queries(
                common_attn_metadata.query_start_loc_cpu,
                common_attn_metadata.seq_lens_cpu_upper_bound,
                max_query_len,
                num_actual_tokens,
            )

        # After DCP distribution, the maximum number of tokens for any rank is
        # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
        # and I is cp_kv_cache_interleave_size.
        # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
        if skip_dcp_context_attention:
            max_dcp_context_kv_len = 0
            scheduler_metadata = None
        else:
            num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
            max_dcp_context_kv_len = (
                (max_seq_len + num_partitions - 1) // num_partitions
            ) * self.cp_kv_cache_interleave_size

            scheduler_metadata = self._get_scheduler_metadata(
                aot_schedule=aot_schedule,
                batch_size=num_reqs,
                cu_query_lens=query_start_loc,
                max_query_len=max_query_len,
                seqlens=dcp_context_kv_lens,
                max_seq_len=max_dcp_context_kv_len,
                causal=False,
                max_num_splits=max_num_splits,
            )
    elif use_cascade:
        cu_prefix_query_lens = torch.tensor(
            [0, num_actual_tokens], dtype=torch.int32, device=self.device
        )
        prefix_kv_lens = torch.tensor(
            [common_prefix_len], dtype=torch.int32, device=self.device
        )
        # Use GPU tensor directly - no CPU sync needed
        suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len
        prefix_scheduler_metadata = self._get_scheduler_metadata(
            aot_schedule=aot_schedule,
            batch_size=1,
            cu_query_lens=cu_prefix_query_lens,
            max_query_len=num_actual_tokens,
            seqlens=prefix_kv_lens,
            max_seq_len=common_prefix_len,
            causal=False,
            max_num_splits=max_num_splits,
        )
        scheduler_metadata = self._get_scheduler_metadata(
            aot_schedule=aot_schedule,
            batch_size=num_reqs,
            cu_query_lens=query_start_loc,
            max_query_len=max_query_len,
            seqlens=suffix_kv_lens,
            max_seq_len=max_seq_len - common_prefix_len,
            causal=True,
            max_num_splits=max_num_splits,
        )
    else:
        scheduler_metadata = self._get_scheduler_metadata(
            aot_schedule=aot_schedule,
            batch_size=num_reqs,
            cu_query_lens=query_start_loc,
            max_query_len=max_query_len,
            seqlens=seq_lens,
            max_seq_len=max_seq_len,
            causal=causal,
            max_num_splits=max_num_splits,
        )
    scheduler_metadata = self._store_scheduler_metadata(scheduler_metadata)

    if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
        raise ValueError(
            f"Per-request causal tensor must be int32, got {causal.dtype}. "
            "Casting here would allocate a fresh tensor each build() and "
            "break FULL CUDA graph replay; allocate the buffer as int32."
        )

    attn_metadata = FlashAttentionMetadata(
        num_actual_tokens=num_actual_tokens,
        max_query_len=max_query_len,
        query_start_loc=query_start_loc,
        max_seq_len=max_seq_len,
        seq_lens=seq_lens,
        block_table=block_table_tensor,
        slot_mapping=slot_mapping,
        max_dcp_context_kv_len=max_dcp_context_kv_len,
        dcp_context_kv_lens=dcp_context_kv_lens,
        num_decode_reqs=num_decode_reqs,
        num_prefill_reqs=num_prefill_reqs,
        num_decode_tokens=num_decode_tokens,
        num_prefill_tokens=num_prefill_tokens,
        use_cascade=use_cascade,
        common_prefix_len=common_prefix_len,
        scheduler_metadata=scheduler_metadata,
        cu_prefix_query_lens=cu_prefix_query_lens,
        prefix_kv_lens=prefix_kv_lens,
        suffix_kv_lens=suffix_kv_lens,
        prefix_scheduler_metadata=prefix_scheduler_metadata,
        max_num_splits=max_num_splits,
        causal=causal,
    )

    # Compute mm_prefix range tensor if the batch contains
    # multimodal tokens with bidirectional ranges.  Built for every FA
    # group; Gemma4 nulls the field for its non-sliding layers.
    mm_ranges = common_attn_metadata.mm_req_doc_ranges
    if mm_ranges is not None and self.mm_prefix_query_ranges_np is not None:
        # The upper bound is exact for prefill rows, which is where
        # mm_prefix ranges live; decode rows only ever get an optimistic
        # (larger) context, moving them further past every range.
        assert common_attn_metadata.seq_lens_cpu_upper_bound is not None, (
            "mm_prefix requires seq_lens_cpu_upper_bound"
        )
        num_mm_tokens = fill_mm_prefix_query_ranges(
            self.mm_prefix_query_ranges_np,
            mm_ranges,
            common_attn_metadata.query_start_loc_cpu,
            common_attn_metadata.seq_lens_cpu_upper_bound,
        )
        if num_mm_tokens > 0:
            assert self.mm_prefix_query_ranges_cpu is not None
            assert self.mm_prefix_query_ranges_gpu is not None
            mm_query_ranges = self.mm_prefix_query_ranges_gpu[:num_mm_tokens]
            mm_query_ranges.copy_(
                self.mm_prefix_query_ranges_cpu[:num_mm_tokens],
                non_blocking=True,
            )
            attn_metadata.mm_prefix_query_range_tensor = mm_query_ranges

    # R-SWA: copy prefix lengths into persistent buffers (outside the
    # compiled region) so forward() never allocates during CUDA graph
    # capture.  rswa_window is a static model config scalar read here.
    if (
        self.rswa_window is not None
        and common_attn_metadata.rswa_prefix_lens is not None
    ):
        assert self.persistent_rswa_prefix_lens is not None
        assert self.persistent_rswa_window_tensor is not None
        src = common_attn_metadata.rswa_prefix_lens
        rswa_prefix_lens = self.persistent_rswa_prefix_lens[:num_reqs]
        rswa_prefix_lens.copy_(src[:num_reqs], non_blocking=True)
        attn_metadata.rswa_prefix_lens = rswa_prefix_lens
        attn_metadata.rswa_window = self.rswa_window
        attn_metadata.rswa_window_tensor = self.persistent_rswa_window_tensor

    return attn_metadata

_get_sliding_window_configs(vllm_config)

Get the set of all sliding window configs used in the model.

Only inspects FlashAttentionImpl layers, including implementations wrapped by a composite backend. Other backends (e.g. TurboQuant, MLA) use their own metadata builders and are skipped.

Source code in vllm/v1/attention/backends/flash_attn.py

def _get_sliding_window_configs(
    vllm_config: VllmConfig,
) -> set[tuple[int, int] | None]:
    """Get the set of all sliding window configs used in the model.

    Only inspects FlashAttentionImpl layers, including implementations wrapped
    by a composite backend. Other backends (e.g. TurboQuant, MLA) use their own
    metadata builders and are skipped.
    """
    sliding_window_configs: set[tuple[int, int] | None] = set()
    layers = get_layers_from_vllm_config(vllm_config, Attention)
    for layer in layers.values():
        for impl in layer.impl.get_impl_variants():
            if isinstance(impl, FlashAttentionImpl):
                sliding_window_configs.add(impl.sliding_window)
    return sliding_window_configs

_make_mm_prefix_mask_mod(sliding_window=0, sliding_window_left=None) cached

Build a CuTE-DSL mask_mod implementing (causal AND sliding_window) OR mm_prefix.

Cached so identical (sliding_window, sliding_window_left) reuse the same function object. FA4's hash_callable mixes repr() of closure cells into the compile key; the nested _load_q_range would otherwise get a new address each call and force a full JIT recompile every forward.

The FA4 kernel passes local q_idx (0-based within the current prefill chunk) while kv_idx is absolute (0-based over the full KV cache). We recover the absolute Q position via q_abs = q_idx + seqlen_k - seqlen_q (the context-length offset) so that causal, sliding-window, and mm_prefix range comparisons all use consistent absolute positions. This matches the Triton reference path (compute_kv_seq_mask).

aux_tensors[0] holds the absolute [start, end] bounds of the mm_prefix range containing each scheduled query token ((-1, -1) when none), and aux_tensors[1] is cu_seqlens_q, used to turn the local q_idx into a packed row index. Because mm_prefix ranges never overlap, r_start <= kv_idx <= r_end is exactly "query and key share a range", so no key-side lookup is needed and kv_idx is never used as an index. The (-1, -1) sentinel falls out for free: kv_idx <= -1 is false for every valid key.

sliding_window_left enforces the sliding window on the causal term (None = full causal, no window). sliding_window clamps the bidirectional block to the window (0 = unclamped; >0 = Gemma4 local layers via mm_prefix_clamp_sliding_window).

Source code in vllm/v1/attention/backends/flash_attn.py

@functools.cache
def _make_mm_prefix_mask_mod(
    sliding_window: int = 0,
    sliding_window_left: int | None = None,
):
    """Build a CuTE-DSL mask_mod implementing
    ``(causal AND sliding_window) OR mm_prefix``.

    Cached so identical ``(sliding_window, sliding_window_left)`` reuse the
    same function object. FA4's ``hash_callable`` mixes ``repr()`` of closure
    cells into the compile key; the nested ``_load_q_range`` would otherwise
    get a new address each call and force a full JIT recompile every forward.

    The FA4 kernel passes *local* ``q_idx`` (0-based within the current
    prefill chunk) while ``kv_idx`` is absolute (0-based over the full
    KV cache).  We recover the absolute Q position via
    ``q_abs = q_idx + seqlen_k - seqlen_q`` (the context-length offset)
    so that causal, sliding-window, and mm_prefix range comparisons all
    use consistent absolute positions.  This matches the Triton
    reference path (``compute_kv_seq_mask``).

    ``aux_tensors[0]`` holds the absolute ``[start, end]`` bounds of the
    mm_prefix range containing each scheduled query token (``(-1, -1)`` when
    none), and ``aux_tensors[1]`` is ``cu_seqlens_q``, used to turn the local
    ``q_idx`` into a packed row index.  Because mm_prefix ranges never overlap,
    ``r_start <= kv_idx <= r_end`` is exactly "query and key share a range", so
    no key-side lookup is needed and ``kv_idx`` is never used as an index.
    The ``(-1, -1)`` sentinel falls out for free: ``kv_idx <= -1`` is false for
    every valid key.

    ``sliding_window_left`` enforces the sliding window on the causal
    term (None = full causal, no window).  ``sliding_window`` clamps the
    bidirectional block to the window (0 = unclamped; >0 = Gemma4 local
    layers via ``mm_prefix_clamp_sliding_window``).
    """
    import cutlass
    import cutlass.cute as cute
    from cutlass import Int32  # type: ignore[attr-defined]

    from vllm.vllm_flash_attn.cute.utils import (  # type: ignore[import-untyped]
        scalar_to_ssa,
        ssa_to_scalar,
    )

    @cute.jit
    def _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx):
        """Load the mm_prefix range bounds for this query row.

        Both loads depend only on the query index, so they hoist out of the
        unrolled per-element mask loop.  ``ssa_to_scalar`` reads lane 0, so one
        call must not span query rows (hence the ``__vec_size__`` pin below).
        Clamping keeps ``token_idx`` in bounds for partial tiles and padded
        (``seqlen_q == 0``) rows; neither produces output.
        """
        q_ranges = aux_tensors[0]
        cu_seqlens_q = aux_tensors[1]
        b = batch_idx[0]
        q_local = cutlass.min(ssa_to_scalar(q_idx), seqlen_info.seqlen_q - Int32(1))
        token_idx = cutlass.max(cu_seqlens_q[b] + q_local, Int32(0))
        return (
            scalar_to_ssa(q_ranges[token_idx, 0], Int32),
            scalar_to_ssa(q_ranges[token_idx, 1], Int32),
        )

    if sliding_window_left is not None:

        @cute.jit
        def mm_prefix_mask_mod(
            batch_idx: cute.TensorSSA,
            head_idx: cute.TensorSSA,
            q_idx: cute.TensorSSA,
            kv_idx: cute.TensorSSA,
            seqlen_info,
            aux_tensors,
        ):
            ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
            q_abs = q_idx + ctx_off
            sw = scalar_to_ssa(Int32(sliding_window_left), Int32)
            keep = (kv_idx <= q_abs) & ((q_abs - kv_idx) < sw)
            r_start, r_end = _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx)
            mm = (kv_idx >= r_start) & (kv_idx <= r_end)
            if sliding_window > 0:
                mm = mm & ((q_abs - kv_idx) < sw)
            keep = keep | mm
            return keep

    else:

        @cute.jit
        def mm_prefix_mask_mod(
            batch_idx: cute.TensorSSA,
            head_idx: cute.TensorSSA,
            q_idx: cute.TensorSSA,
            kv_idx: cute.TensorSSA,
            seqlen_info,
            aux_tensors,
        ):
            ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
            q_abs = q_idx + ctx_off
            keep = kv_idx <= q_abs
            r_start, r_end = _load_q_range(q_idx, seqlen_info, aux_tensors, batch_idx)
            keep = keep | ((kv_idx >= r_start) & (kv_idx <= r_end))
            return keep

    mm_prefix_mask_mod.use_fast_sampling = True
    mm_prefix_mask_mod.__vec_size__ = 1  # _load_q_range takes lane 0 of q_idx
    return mm_prefix_mask_mod

_make_rswa_mask_mod()

Build a CuTE-DSL mask_mod for Reference Sliding Window Attention (R-SWA).

FA4 varlen + paged-KV convention (verified from cute/mask.py apply_mask): q_idx = LOCAL query-token offset (0 .. seqlen_q - 1) within this sequence. kv_idx = LOCAL KV-token position (0 .. seqlen_k - 1) within this sequence.

To recover the ABSOLUTE token position (needed for causal and the sliding window distance), use the standard offset: abs_q = q_idx + (seqlen_k - seqlen_q)

R-SWA keep condition

abs_q >= kv_idx (causal: KV at or before the query) AND (kv_idx < prefix_len (global prefix is always visible) OR abs_q - kv_idx < window) (generated tokens: sliding window)

aux_tensors[0]: prefix_lens [num_reqs] int32 — per-request prefill length. aux_tensors[1]: rswa_window [1] int32 — decode sliding window size.

use_fast_sampling=True lets FA4 skip fully-masked KV blocks (gap blocks) without loading their data.

Source code in vllm/v1/attention/backends/flash_attn.py

def _make_rswa_mask_mod():
    """Build a CuTE-DSL mask_mod for Reference Sliding Window Attention (R-SWA).

    FA4 varlen + paged-KV convention (verified from cute/mask.py apply_mask):
      q_idx  = LOCAL query-token offset (0 .. seqlen_q - 1) within this sequence.
      kv_idx = LOCAL KV-token position (0 .. seqlen_k - 1) within this sequence.

    To recover the ABSOLUTE token position (needed for causal and the sliding
    window distance), use the standard offset:
      abs_q = q_idx + (seqlen_k - seqlen_q)

    R-SWA keep condition:
      abs_q >= kv_idx                    (causal: KV at or before the query)
      AND (kv_idx < prefix_len           (global prefix is always visible)
           OR  abs_q - kv_idx < window)  (generated tokens: sliding window)

    aux_tensors[0]: prefix_lens [num_reqs] int32 — per-request prefill length.
    aux_tensors[1]: rswa_window [1]        int32 — decode sliding window size.

    use_fast_sampling=True lets FA4 skip fully-masked KV blocks (gap blocks)
    without loading their data.
    """
    import cutlass.cute as cute
    from cutlass import Int32  # type: ignore[attr-defined]

    from vllm.vllm_flash_attn.cute.utils import (  # type: ignore[import-untyped]
        scalar_to_ssa,
    )

    @cute.jit
    def rswa_mask_mod(
        batch_idx: cute.TensorSSA,
        head_idx: cute.TensorSSA,
        q_idx: cute.TensorSSA,
        kv_idx: cute.TensorSSA,
        seqlen_info,
        aux_tensors,
    ):
        b = batch_idx[0]
        prefix_len = scalar_to_ssa(aux_tensors[0][b], Int32)
        window = scalar_to_ssa(aux_tensors[1][0], Int32)
        # Convert local q offset to absolute token position.
        offset = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
        abs_q = q_idx + offset
        causal = kv_idx <= abs_q
        in_prefix = kv_idx < prefix_len
        in_window = (abs_q - kv_idx) < window
        return causal & (in_prefix | in_window)

    rswa_mask_mod.use_fast_sampling = True
    return rswa_mask_mod

_maybe_symmetrize_window(window, causal)

Make a causal sliding window (w, 0) symmetric (w, w) when attention is non-causal, so bidirectional queries attend in both directions. Leaves full-attention (-1, -1) and already-symmetric windows untouched.

Source code in vllm/v1/attention/backends/flash_attn.py

def _maybe_symmetrize_window(
    window: tuple[int, int] | None,
    causal: bool | torch.Tensor,
) -> tuple[int, int] | None:
    """Make a causal sliding window ``(w, 0)`` symmetric ``(w, w)`` when attention
    is non-causal, so bidirectional queries attend in both directions. Leaves
    full-attention ``(-1, -1)`` and already-symmetric windows untouched.
    """
    non_causal = isinstance(causal, torch.Tensor) or causal is False
    if window is not None and window[0] >= 0 and window[1] == 0 and non_causal:
        return (window[0], window[0])
    return window

use_cascade_attention(common_prefix_len, query_lens, num_query_heads, num_kv_heads, use_alibi, use_sliding_window, use_local_attention, num_sms, dcp_world_size)

Decide whether to use cascade attention.

This function 1) checks whether cascade attention is supported with the given configuration, and 2) heuristically decides whether using cascade attention can improve performance.

Source code in vllm/v1/attention/backends/flash_attn.py

def use_cascade_attention(
    common_prefix_len: int,
    query_lens: np.ndarray,
    num_query_heads: int,
    num_kv_heads: int,
    use_alibi: bool,
    use_sliding_window: bool,
    use_local_attention: bool,
    num_sms: int,
    dcp_world_size: int,
) -> bool:
    """Decide whether to use cascade attention.

    This function 1) checks whether cascade attention is supported with the
    given configuration, and 2) heuristically decides whether using cascade
    attention can improve performance.
    """
    # Too short common prefix. Probably not worth using cascade attention.
    # We use an arbitrary threshold of 256 tokens. TODO: Tune this threshold.
    # NOTE(woosuk): This is the common case. We should return False as soon as
    # possible to avoid any unnecessary computation.
    if common_prefix_len < 256:
        return False
    # Cascade attention is currently not supported with these variants.
    if use_alibi or use_sliding_window or use_local_attention:
        return False
    # Too few queries. Probably not worth using cascade attention.
    # We use an arbitrary threshold of 8 queries. TODO: Tune this threshold.
    num_reqs = len(query_lens)
    if num_reqs < 8:
        return False
    # disable cascade attention for DCP
    if dcp_world_size > 1:
        return False

    # Heuristics to decide whether using cascade attention is beneficial.
    # 1. When FlashDecoding is not used for normal attention, cascade attention
    #    is likely to be faster since it saves memory bandwidth.
    num_queries_per_kv = num_query_heads // num_kv_heads
    # The criteria for using FlashDecoding can be found in the following link:
    # https://github.com/vllm-project/flash-attention/blob/96266b1111111f3d11aabefaf3bacbab6a89d03c/csrc/flash_attn/flash_api.cpp#L535
    use_flash_decoding = (
        num_queries_per_kv > 1
        and not use_sliding_window
        and not use_alibi
        and np.all(query_lens == 1)
    )
    if not use_flash_decoding:
        # Use cascade attention.
        return True

    # 2. When FlashDecoding is used for normal attention, it is not clear
    #    whether cascade attention is beneficial, because FlashDecoding can
    #    launch more CTAs than cascade attention.
    #    We use a simple performance model to compare the two methods.
    #    NOTE(woosuk): The performance model is very rough and may not be
    #    accurate.
    num_tokens = num_reqs
    # NOTE(woosuk): These are default tile sizes. flash-attn might use
    # different tile sizes (e.g., 64 or 256) depending on the configuration.
    q_tile_size = 128
    kv_tile_size = 128
    num_prefix_tiles = cdiv(common_prefix_len, kv_tile_size)

    cascade_ctas = num_query_heads * cdiv(num_tokens, q_tile_size)
    cascade_waves = cdiv(cascade_ctas, num_sms)
    cascade_time = cascade_waves * num_prefix_tiles

    flash_decoding_ctas = (
        num_reqs * num_kv_heads * cdiv(num_queries_per_kv, q_tile_size)
    )
    flash_decoding_ctas *= num_prefix_tiles
    flash_decoding_time = cdiv(flash_decoding_ctas, num_sms)

    # Use cascade attention if it is faster than FlashDecoding.
    return cascade_time < flash_decoding_time