sglang 0.4.5__py3-none-any.whl → 0.4.5.post1__py3-none-any.whl

sglang/srt/layers/attention/flashattention_backend.py

@@ -1,51 +1,56 @@
 from __future__ import annotations
 
-import numpy as np
-
-from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
-
-"""
-Support different attention backends.
-Now there are three backends: FlashInfer, Triton and FlashAttention.
-Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
-"""
-
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Optional, Union
 
+import numpy as np
 import torch
 
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
 
-from sgl_kernel.flash_attn import flash_attn_with_kvcache
+from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
 
 
 @dataclass
 class FlashAttentionMetadata:
     """Metadata to be init once in the model forward pass,
-    each layer's forward pass can reuse the metadata."""
+    each layer's forward pass can reuse the metadata.
 
-    # Cumulative sequence lengths for query
-    cu_seqlens_q: torch.Tensor = None
-    # Cumulative sequence lengths for key
-    cu_seqlens_k: torch.Tensor = None
+    For each init metadata function, we will try set up them in below order
+    """
+
+    # Sequence lengths for the forward batch
+    cache_seqlens_int32: torch.Tensor = None
     # Maximum sequence length for query
     max_seq_len_q: int = 0
     # Maximum sequence length for key
     max_seq_len_k: int = 0
+    # Cumulative sequence lengths for query
+    cu_seqlens_q: torch.Tensor = None
+    # Cumulative sequence lengths for key
+    cu_seqlens_k: torch.Tensor = None
     # Window size (typically used by Gemma)
     window_size: tuple = (-1, -1)
     # Page table, the index of KV Cache Tables/Blocks
     page_table: torch.Tensor = None
+
+    # Encoder metadata
+    # Cumulative sequence lengths for encoder key
+    encoder_cu_seqlens_k: torch.Tensor = None
+    # Maximum sequence length for encoder key
+    encoder_max_seq_len_k: int = 0
     # Sequence lengths for the forward batch
-    cache_seqlens_int32: torch.Tensor = None
+    encoder_lens_int32: torch.Tensor = None
+    # Page table for the encoder
+    encoder_page_table: torch.Tensor = None
 
     @dataclass
     class LocalAttentionMetadata:
@@ -231,7 +236,11 @@ def make_local_attention_virtual_batches(
         np.arange(pages_per_local_batch, dtype=np.int32),
         (virtual_batches, pages_per_local_batch),
     ) + np.expand_dims(block_starts, axis=1)
-    block_indices = block_indices.flatten()
+    # Ensure block_indices doesn't exceed block_table dimensions
+    # This is a critical safety check that prevents index out of bounds errors
+    # when dealing with large sequences (>8192 tokens) or when the block_table
+    # dimensions are smaller than what would be needed for the full attention chunk size.
+    block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1)
     batch_indices = np.repeat(
         np.arange(actual_batch_size, dtype=np.int32),
         local_blocks * pages_per_local_batch,
@@ -270,9 +279,9 @@ class FlashAttentionBackend(AttentionBackend):
         self,
         model_runner: ModelRunner,
         skip_prefill: bool = False,
+        speculative_step_id=0,
         topk=0,
         speculative_num_steps=0,
-        step_id=0,
     ):
         super().__init__()
 
@@ -287,20 +296,20 @@ class FlashAttentionBackend(AttentionBackend):
         self.decode_cuda_graph_metadata = {}
         self.target_verify_metadata = {}
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.kv_cache_dtype = model_runner.kv_cache_dtype
+        self.kv_cache_dtype_str = model_runner.server_args.kv_cache_dtype
         self.page_size = model_runner.page_size
         self.use_mla = (
             model_runner.model_config.attention_arch == AttentionArch.MLA
         ) and (not global_server_args_dict["disable_mla"])
         self.skip_prefill = skip_prefill
 
-        # TODO: Support Topk > 1 for FlashAttentionBackend Spec Decoding
-        assert (
-            topk <= 1
-        ), "topk must be 1 (if spec decoding) or 0 (if no spec decoding) for FlashAttentionBackend"
-
-        self.topk = 1
-        self.step_id = step_id
+        self.topk = topk
         self.speculative_num_steps = speculative_num_steps
+        self.speculative_num_draft_tokens = (
+            model_runner.server_args.speculative_num_draft_tokens
+        )
+        self.speculative_step_id = speculative_step_id
 
         # Local attention settings
         self.attention_chunk_size = (
@@ -310,71 +319,59 @@ class FlashAttentionBackend(AttentionBackend):
         )
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
-        """Initialize forward metadata to cache repetitive calculations."""
+        """Initialize forward metadata hence all layers in the forward pass can reuse it."""
         metadata = FlashAttentionMetadata()
         seqlens_in_batch = forward_batch.seq_lens
         batch_size = len(seqlens_in_batch)
         device = seqlens_in_batch.device
-        if forward_batch.forward_mode.is_decode():
-            # Skip Prefill or Draft Decode
-            # Note: Draft Decode will be ran on the Draft Worker
+
+        if forward_batch.forward_mode.is_decode_or_idle():
+            # Draft Decode
             if forward_batch.spec_info is not None:
+                metadata.cache_seqlens_int32 = (
+                    seqlens_in_batch + (self.speculative_step_id + 1)
+                ).to(torch.int32)
+                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + (
+                    self.speculative_step_id + 1
+                )
                 metadata.cu_seqlens_q = torch.arange(
                     0, batch_size + 1, dtype=torch.int32, device=device
                 )
-                seq_lens_with_decode = seqlens_in_batch + (self.step_id + 1)
-                metadata.cache_seqlens_int32 = seq_lens_with_decode.to(torch.int32)
                 metadata.cu_seqlens_k = torch.nn.functional.pad(
                     torch.cumsum(
                         metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
                     ),
                     (1, 0),
                 )
-                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + (
-                    self.step_id + 1
-                )
                 metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
                     forward_batch.req_pool_indices, : metadata.max_seq_len_k
                 ]
-                cache_loc = forward_batch.out_cache_loc.view(
-                    self.speculative_num_steps, -1
-                ).T
-
-                for idx, single_seq_len in enumerate(seq_lens_with_decode):
-                    real_bsz_start_idx = idx
-                    real_bsz_end_idx = idx + 1
-                    metadata.page_table[
-                        real_bsz_start_idx:real_bsz_end_idx,
-                        (single_seq_len - (self.step_id + 1)) : single_seq_len,
-                    ] = cache_loc[
-                        real_bsz_start_idx:real_bsz_end_idx, : (self.step_id + 1)
-                    ]
-            else:  # Normal Decode without Spec Decoding
+            else:
+                # Normal Decode
                 metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
+                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
+                metadata.cu_seqlens_q = torch.arange(
+                    0, batch_size + 1, dtype=torch.int32, device=device
+                )
                 metadata.cu_seqlens_k = torch.nn.functional.pad(
                     torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
                 )
-                metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
                 metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
                     forward_batch.req_pool_indices, : metadata.max_seq_len_k
                 ]
-                metadata.cu_seqlens_q = torch.arange(
-                    0, batch_size + 1, dtype=torch.int32, device=device
-                )
         elif forward_batch.forward_mode.is_target_verify():
-            # Note: Target Verify will be ran on the Target Worker
-            draft_token_num = forward_batch.spec_info.draft_token_num
             metadata.cache_seqlens_int32 = (
-                forward_batch.seq_lens + draft_token_num
+                forward_batch.seq_lens + self.speculative_num_draft_tokens
             ).to(torch.int32)
-            metadata.max_seq_len_q = draft_token_num
+            metadata.max_seq_len_q = self.speculative_num_draft_tokens
             metadata.max_seq_len_k = (
-                forward_batch.seq_lens_cpu.max().item() + draft_token_num
+                forward_batch.seq_lens_cpu.max().item()
+                + self.speculative_num_draft_tokens
             )
             metadata.cu_seqlens_q = torch.arange(
                 0,
-                batch_size * draft_token_num + 1,
-                draft_token_num,
+                batch_size * self.speculative_num_draft_tokens + 1,
+                self.speculative_num_draft_tokens,
                 dtype=torch.int32,
                 device=device,
             )
@@ -386,32 +383,28 @@ class FlashAttentionBackend(AttentionBackend):
                 forward_batch.req_pool_indices, : metadata.max_seq_len_k
             ]
 
-        elif forward_batch.forward_mode.is_extend_or_draft_extend():
-            # Normal or Draft Extend (Both of them will be ran on the Target Worker)
+        elif forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed():
             metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
+            metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
             metadata.cu_seqlens_k = torch.nn.functional.pad(
                 torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
             )
-            # Precompute maximum sequence length
-            metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
-            # Precompute page table
             metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
                 forward_batch.req_pool_indices, : metadata.max_seq_len_k
             ]
 
-            # Precompute cumulative sequence lengths
             if (
                 any(forward_batch.extend_prefix_lens_cpu)
                 or forward_batch.forward_mode == ForwardMode.DRAFT_EXTEND
             ):
                 extend_seq_lens = forward_batch.extend_seq_lens
+                metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
                 metadata.cu_seqlens_q = torch.nn.functional.pad(
                     torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0)
                 )
-                metadata.max_seq_len_q = max(forward_batch.extend_seq_lens_cpu)
             else:
-                metadata.cu_seqlens_q = metadata.cu_seqlens_k
                 metadata.max_seq_len_q = metadata.max_seq_len_k
+                metadata.cu_seqlens_q = metadata.cu_seqlens_k
 
             # Setup local attention if enabled
             if (
@@ -458,7 +451,31 @@ class FlashAttentionBackend(AttentionBackend):
                 )
                 metadata.local_attn_metadata = local_metadata
 
-        # Precompute strided indices
+        # Encoder metadata for cross attention
+        if forward_batch.encoder_lens is not None:
+            assert (
+                forward_batch.encoder_lens.numel() == 1
+            ), "Only encoder size 1 is supported for now"
+
+            metadata.encoder_lens_int32 = forward_batch.encoder_lens.to(torch.int32)
+            metadata.encoder_cu_seqlens_k = torch.nn.functional.pad(
+                torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
+                (1, 0),
+            )
+            metadata.encoder_max_seq_len_k = metadata.encoder_lens_int32.max().item()
+            metadata.encoder_page_table = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices, : metadata.encoder_max_seq_len_k
+            ]
+
+            # Currently only support forward_batch.encoder_lens.numel() == 1
+            metadata.page_table = forward_batch.req_to_token_pool.req_to_token[
+                forward_batch.req_pool_indices,
+                metadata.encoder_max_seq_len_k : (
+                    metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
+                ),
+            ]
+
+        # Convert the page table to a strided format which is needed by FA3 API
         if self.page_size > 1:
             self.strided_indices = torch.arange(
                 0, metadata.page_table.shape[1], self.page_size, device=self.device
@@ -498,7 +515,7 @@ class FlashAttentionBackend(AttentionBackend):
                         v,
                     )
 
-        # Use precomputed metadata
+        # Use precomputed metadata across all layers
         metadata = self.forward_metadata
 
         # Calculate window size (can be moved to metadata if layer properties don't change)
@@ -506,9 +523,18 @@ class FlashAttentionBackend(AttentionBackend):
         # here is two side inclusive
         window_size = (
             (layer.sliding_window_size, 0)
-            if layer.sliding_window_size is not None
+            if layer.sliding_window_size is not None and layer.sliding_window_size > -1
             else (-1, -1)
         )
+        k_descale, v_descale = None, None
+        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+        # has corresponding quantization method so that layer.k_scale is not None
+        if self.kv_cache_dtype_str != "auto" and layer.k_scale is not None:
+            descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+            k_descale = layer.k_scale.expand(descale_shape)
+            v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
+        causal = not layer.is_cross_attention
 
         # Check if we should use local attention
         use_local_attn = (
@@ -536,14 +562,21 @@ class FlashAttentionBackend(AttentionBackend):
         # Use Flash Attention for prefill
         if not self.use_mla:
             # Do multi-head attention
-            kv_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-            key_cache, value_cache = kv_cache[0], kv_cache[1]
+            key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
+                layer.layer_id
+            )
             key_cache = key_cache.view(
                 -1, self.page_size, layer.tp_k_head_num, layer.head_dim
             )
             value_cache = value_cache.view(
                 -1, self.page_size, layer.tp_v_head_num, layer.head_dim
             )
+            if layer.is_cross_attention:
+                page_table = metadata.encoder_page_table
+                cache_seqlens = metadata.encoder_lens_int32
+                cu_seqlens_k = metadata.encoder_cu_seqlens_k
+                window_size = (-1, -1)
+
             o = flash_attn_with_kvcache(
                 q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                 k_cache=key_cache,
@@ -554,48 +587,93 @@ class FlashAttentionBackend(AttentionBackend):
                 cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
                 max_seqlen_q=max_seqlen_q,
                 softmax_scale=layer.scaling,
-                causal=True,
+                causal=causal,
                 window_size=window_size,
                 softcap=layer.logit_cap,
-                k_descale=layer.k_scale,
-                v_descale=layer.v_scale,
+                k_descale=k_descale,
+                v_descale=v_descale,
             )
+            return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
         else:
-            # Do absorbed multi-latent attention
-            kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-            k_rope = kv_cache[:, :, layer.v_head_dim :]
-            c_kv = kv_cache[:, :, : layer.v_head_dim]
-            k_rope_cache = k_rope.view(
-                -1,
-                self.page_size,
-                layer.tp_k_head_num,
-                layer.head_dim - layer.v_head_dim,
-            )
-            c_kv_cache = c_kv.view(
-                -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
-            )
-
-            q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
-            q_nope = q_all[:, :, : layer.v_head_dim]
-            q_rope = q_all[:, :, layer.v_head_dim :]
-            o = flash_attn_with_kvcache(
-                q=q_rope,
-                k_cache=k_rope_cache,
-                v_cache=c_kv_cache,
-                qv=q_nope,
-                page_table=page_table,
-                cache_seqlens=cache_seqlens,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
-                max_seqlen_q=max_seqlen_q,
-                softmax_scale=layer.scaling,
-                causal=True,
-                softcap=layer.logit_cap,
-                k_descale=layer.k_scale,
-                v_descale=layer.v_scale,
-            )
+            if (
+                not global_server_args_dict["disable_chunked_prefix_cache"]
+                and forward_batch.attn_attend_prefix_cache is not None
+                and not forward_batch.forward_mode.is_target_verify()
+                and not forward_batch.forward_mode.is_draft_extend()
+            ):
+                # Do multi-head attention with chunked prefix cache
+
+                if forward_batch.attn_attend_prefix_cache:
+                    # MHA for chunked prefix kv cache when running model with MLA
+                    assert forward_batch.prefix_chunk_idx is not None
+                    assert forward_batch.prefix_chunk_cu_seq_lens is not None
+                    assert forward_batch.prefix_chunk_max_seq_lens is not None
+
+                    chunk_idx = forward_batch.prefix_chunk_idx
+                    assert chunk_idx >= 0
+
+                    output, lse, *rest = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=forward_batch.prefix_chunk_cu_seq_lens[chunk_idx],
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=forward_batch.prefix_chunk_max_seq_lens[chunk_idx],
+                        softmax_scale=layer.scaling,
+                        causal=False,
+                        return_softmax_lse=True,
+                    )
+                else:
+                    # MHA for extend part of sequence without attending prefix kv cache
+                    output, lse, *rest = flash_attn_varlen_func(
+                        q=q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                        k=k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                        v=v.view(-1, layer.tp_k_head_num, layer.v_head_dim),
+                        cu_seqlens_q=metadata.cu_seqlens_q,
+                        cu_seqlens_k=metadata.cu_seqlens_q,
+                        max_seqlen_q=metadata.max_seq_len_q,
+                        max_seqlen_k=metadata.max_seq_len_q,
+                        softmax_scale=layer.scaling,
+                        causal=True,
+                        return_softmax_lse=True,
+                    )
+                return output, lse
+            else:
+                # Do absorbed multi-latent attention
+                kv_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+                k_rope = kv_cache[:, :, layer.v_head_dim :]
+                c_kv = kv_cache[:, :, : layer.v_head_dim]
+                k_rope_cache = k_rope.view(
+                    -1,
+                    self.page_size,
+                    layer.tp_k_head_num,
+                    layer.head_dim - layer.v_head_dim,
+                )
+                c_kv_cache = c_kv.view(
+                    -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
+                )
+                q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                q_nope = q_all[:, :, : layer.v_head_dim]
+                q_rope = q_all[:, :, layer.v_head_dim :]
+                o = flash_attn_with_kvcache(
+                    q=q_rope,
+                    k_cache=k_rope_cache,
+                    v_cache=c_kv_cache,
+                    qv=q_nope,
+                    page_table=page_table,
+                    cache_seqlens=cache_seqlens,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
+                    max_seqlen_q=max_seqlen_q,
+                    softmax_scale=layer.scaling,
+                    causal=True,
+                    softcap=layer.logit_cap,
+                    k_descale=k_descale,
+                    v_descale=v_descale,
+                )
 
-        return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+                return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
     def forward_decode(
         self,
@@ -606,8 +684,6 @@ class FlashAttentionBackend(AttentionBackend):
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ) -> torch.Tensor:
-        """Forward pass with FlashAttention using precomputed metadata."""
-        # Save KV cache if needed
         if k is not None:
             assert v is not None
             if save_kv_cache:
@@ -628,7 +704,7 @@ class FlashAttentionBackend(AttentionBackend):
                         v,
                     )
 
-        # Use precomputed metadata
+        # Use precomputed metadata across all layers
         metadata = self.forward_metadata
 
         # Calculate window size (can be moved to metadata if layer properties don't change)
@@ -636,17 +712,27 @@ class FlashAttentionBackend(AttentionBackend):
         # here is two side inclusive
         window_size = (
             (layer.sliding_window_size, 0)
-            if layer.sliding_window_size is not None
+            if layer.sliding_window_size is not None and layer.sliding_window_size > -1
             else (-1, -1)
         )
-        page_table = metadata.page_table
+        causal = not layer.is_cross_attention
+
+        k_descale, v_descale = None, None
+        # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention
+        # has corresponding quantization method so that layer.k_scale is not None
+        if self.kv_cache_dtype_str != "auto":
+            if layer.k_scale is not None:
+                descale_shape = (forward_batch.batch_size, layer.tp_k_head_num)
+                k_descale = layer.k_scale.expand(descale_shape)
+                v_descale = layer.v_scale.expand(descale_shape)
+            q = q.to(self.kv_cache_dtype)
 
         if not self.use_mla:
             # Do multi-head attention
 
-            # Get KV cache
-            kv_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-            key_cache, value_cache = kv_cache[0], kv_cache[1]
+            key_cache, value_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
+                layer.layer_id
+            )
             key_cache = key_cache.view(
                 -1, self.page_size, layer.tp_k_head_num, layer.head_dim
             )
@@ -654,23 +740,32 @@ class FlashAttentionBackend(AttentionBackend):
                 -1, self.page_size, layer.tp_v_head_num, layer.head_dim
             )
 
-            # Pre-reshape query tensor
             q_reshaped = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+            if layer.is_cross_attention:
+                page_table = metadata.encoder_page_table
+                cache_seqlens = metadata.encoder_lens_int32
+                cu_seqlens_k = metadata.encoder_cu_seqlens_k
+                window_size = (-1, -1)
+            else:
+                page_table = metadata.page_table
+                cache_seqlens = metadata.cache_seqlens_int32
+                cu_seqlens_k = metadata.cu_seqlens_k
+
             o = flash_attn_with_kvcache(
                 q=q_reshaped,
                 k_cache=key_cache,
                 v_cache=value_cache,
                 page_table=page_table,
-                cache_seqlens=metadata.cache_seqlens_int32,
+                cache_seqlens=cache_seqlens,
                 cu_seqlens_q=metadata.cu_seqlens_q,
-                cu_seqlens_k_new=metadata.cu_seqlens_k,
+                cu_seqlens_k_new=cu_seqlens_k,
                 max_seqlen_q=1,
                 softmax_scale=layer.scaling,
-                causal=True,
+                causal=causal,
                 window_size=window_size,
                 softcap=layer.logit_cap,
-                k_descale=layer.k_scale,
-                v_descale=layer.v_scale,
+                k_descale=k_descale,
+                v_descale=v_descale,
             )
         else:
             # Do absorbed multi-latent attention
@@ -696,7 +791,7 @@ class FlashAttentionBackend(AttentionBackend):
                 k_cache=k_rope_cache,
                 v_cache=c_kv_cache,
                 qv=q_nope,
-                page_table=page_table,
+                page_table=metadata.page_table,
                 cache_seqlens=metadata.cache_seqlens_int32,
                 cu_seqlens_q=metadata.cu_seqlens_q,
                 cu_seqlens_k_new=metadata.cu_seqlens_k,
@@ -704,8 +799,8 @@ class FlashAttentionBackend(AttentionBackend):
                 softmax_scale=layer.scaling,
                 causal=True,
                 softcap=layer.logit_cap,
-                k_descale=layer.k_scale,
-                v_descale=layer.v_scale,
+                k_descale=k_descale,
+                v_descale=v_descale,
             )
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 
@@ -719,7 +814,13 @@ class FlashAttentionBackend(AttentionBackend):
         to avoid memory allocations.
         """
         self.decode_cuda_graph_metadata = {
-            # Page table for token mapping (batch_size, max_context_len)
+            "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
+            "cu_seqlens_q": torch.arange(
+                0, max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+            "cu_seqlens_k": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
+            ),
             "page_table": torch.zeros(
                 max_bs,
                 (self.max_context_len + self.page_size - 1) // self.page_size,
@@ -735,35 +836,42 @@ class FlashAttentionBackend(AttentionBackend):
             "strided_indices": torch.arange(
                 0, self.max_context_len, self.page_size, device=self.device
             ),
+        }
+
+        self.target_verify_metadata = {
             "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
-            "cu_seqlens_q": torch.arange(
-                0, max_bs + 128, dtype=torch.int32, device=self.device
+            "cu_seqlens_q": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
             ),
             "cu_seqlens_k": torch.zeros(
-                max_bs + 128, dtype=torch.int32, device=self.device
+                max_bs + 1, dtype=torch.int32, device=self.device
             ),
-        }
-
-        self.target_verify_metadata = {
             "page_table": torch.zeros(
                 max_bs,
                 (self.max_context_len + self.page_size - 1) // self.page_size,
                 dtype=torch.int32,
                 device=self.device,
             ),
-            "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
-            "cu_seqlens_q": torch.zeros(
-                max_bs + 128, dtype=torch.int32, device=self.device
-            ),
-            "cu_seqlens_k": torch.zeros(
-                max_bs + 128, dtype=torch.int32, device=self.device
-            ),
-            "max_seqlen_q": 0,
             "strided_indices": torch.arange(
                 0, self.max_context_len, self.page_size, device=self.device
             ),
         }
 
+        self.encoder_metadata = {
+            "encoder_page_table": torch.zeros(
+                max_bs,
+                self.max_context_len,
+                dtype=torch.int32,
+                device=self.device,
+            ),
+            "encoder_lens_int32": torch.zeros(
+                max_bs, dtype=torch.int32, device=self.device
+            ),
+            "encoder_cu_seqlens_k": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+        }
+
     def init_forward_metadata_capture_cuda_graph(
         self,
         bs: int,
@@ -777,27 +885,24 @@ class FlashAttentionBackend(AttentionBackend):
         """Initialize forward metadata for capturing CUDA graph."""
         metadata = FlashAttentionMetadata()
         device = seq_lens.device
-        if forward_mode.is_decode():
+        if forward_mode.is_decode_or_idle():
             if spec_info is not None:
                 # Draft Decode
-                metadata.cu_seqlens_q = torch.arange(
-                    0, bs + 1, dtype=torch.int32, device=device
-                )
                 metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[
                     "cache_seqlens"
                 ][:bs]
-
+                metadata.max_seq_len_k = seq_lens.max().item() + (
+                    self.speculative_step_id + 1
+                )
                 metadata.cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][
                     : bs + 1
                 ]
-
                 metadata.cu_seqlens_k = torch.nn.functional.pad(
                     torch.cumsum(
                         metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
                     ),
                     (1, 0),
                 )
-                metadata.max_seq_len_k = seq_lens.max().item() + (self.step_id + 1)
                 metadata.page_table = self.decode_cuda_graph_metadata[
                     "page_table_draft_decode"
                 ][req_pool_indices, :]
@@ -822,43 +927,49 @@ class FlashAttentionBackend(AttentionBackend):
                 )
             self.decode_cuda_graph_metadata[bs] = metadata
         elif forward_mode.is_target_verify():
-            draft_token_num = spec_info.draft_token_num
-
             metadata.cache_seqlens_int32 = self.target_verify_metadata["cache_seqlens"][
                 :bs
             ]
             metadata.cache_seqlens_int32.copy_(
-                (seq_lens + draft_token_num).to(torch.int32)
+                (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
             )
 
-            metadata.max_seq_len_q = draft_token_num
-            metadata.max_seq_len_k = seq_lens.max().item() + draft_token_num
+            metadata.max_seq_len_q = self.speculative_num_draft_tokens
+            metadata.max_seq_len_k = (
+                seq_lens.max().item() + self.speculative_num_draft_tokens
+            )
 
-            metadata.cu_seqlens_q = self.target_verify_metadata["cu_seqlens_q"][
-                torch.arange(
-                    0,
-                    bs * draft_token_num + 1,
-                    draft_token_num,
-                    dtype=torch.int32,
-                    device=device,
-                )
-            ]
-            cu_k = self.target_verify_metadata["cu_seqlens_k"][: (bs + 1)]
-            cu_k.copy_(
-                torch.nn.functional.pad(
-                    torch.cumsum(
-                        metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
-                    ),
-                    (1, 0),
-                )
+            metadata.cu_seqlens_q = torch.arange(
+                0,
+                bs * self.speculative_num_draft_tokens + 1,
+                self.speculative_num_draft_tokens,
+                dtype=torch.int32,
+                device=device,
             )
-            metadata.cu_seqlens_k = cu_k
+
+            metadata.cu_seqlens_k = self.target_verify_metadata["cu_seqlens_k"][
+                : (bs + 1)
+            ]
+
             metadata.page_table = self.target_verify_metadata["page_table"][
                 req_pool_indices, :
             ]
 
             self.target_verify_metadata[bs] = metadata
 
+        if encoder_lens is not None:
+            encoder_bs = encoder_lens.numel()
+            metadata.encoder_lens_int32 = self.encoder_metadata["encoder_lens_int32"][
+                :encoder_bs
+            ]
+            metadata.encoder_cu_seqlens_k = self.encoder_metadata[
+                "encoder_cu_seqlens_k"
+            ][: (encoder_bs + 1)]
+
+            metadata.encoder_page_table = self.encoder_metadata["encoder_page_table"][
+                req_pool_indices, :
+            ]
+
         self.forward_metadata = metadata
 
     def init_forward_metadata_replay_cuda_graph(
@@ -874,24 +985,21 @@ class FlashAttentionBackend(AttentionBackend):
         out_cache_loc: torch.Tensor = None,
     ):
         # """Initialize forward metadata for replaying CUDA graph."""
-        device = seq_lens.device
         seq_lens = seq_lens[:bs]
-        req_pool_indices = req_pool_indices[:bs]
         seq_lens_cpu = seq_lens_cpu[:bs]
-        if forward_mode.is_decode():
+        req_pool_indices = req_pool_indices[:bs]
+        if forward_mode.is_decode_or_idle():
             metadata = self.decode_cuda_graph_metadata[bs]
 
             if spec_info is not None:
                 # Draft Decode
-                max_len = seq_lens_cpu.max().item()
-                metadata.max_seq_len_k = max_len + (self.step_id + 1)
-
                 metadata.cache_seqlens_int32.copy_(
-                    (seq_lens + (self.step_id + 1)).to(torch.int32)
+                    (seq_lens + (self.speculative_step_id + 1)).to(torch.int32)
                 )
 
-                metadata.max_seq_len_k = seq_lens_cpu.max().item() + (self.step_id + 1)
-
+                metadata.max_seq_len_k = seq_lens_cpu.max().item() + (
+                    self.speculative_step_id + 1
+                )
                 metadata.cu_seqlens_k.copy_(
                     torch.nn.functional.pad(
                         torch.cumsum(
@@ -920,31 +1028,24 @@ class FlashAttentionBackend(AttentionBackend):
                     metadata.max_seq_len_k + self.page_size - 1
                 ) // self.page_size
                 page_indices = self.req_to_token[
-                    :,
-                    self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
+                    req_pool_indices[:, None],
+                    self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages][
+                        None, :
+                    ],
                 ]
-                page_indices = page_indices[req_pool_indices] // self.page_size
+                page_indices //= self.page_size
                 metadata.page_table[:, :max_seq_pages].copy_(page_indices)
                 metadata.page_table[:, max_seq_pages:].fill_(0)
 
         elif forward_mode.is_target_verify():
             metadata = self.target_verify_metadata[bs]
-            draft_token_num = spec_info.draft_token_num
-
-            metadata.cu_seqlens_q.copy_(
-                torch.arange(
-                    0,
-                    bs * draft_token_num + 1,
-                    draft_token_num,
-                    dtype=torch.int32,
-                    device=device,
-                )
-            )
             metadata.cache_seqlens_int32.copy_(
-                (seq_lens + draft_token_num).to(torch.int32)
+                (seq_lens + self.speculative_num_draft_tokens).to(torch.int32)
             )
 
-            metadata.max_seq_len_k = seq_lens_cpu.max().item() + draft_token_num
+            metadata.max_seq_len_k = (
+                seq_lens_cpu.max().item() + self.speculative_num_draft_tokens
+            )
             metadata.cu_seqlens_k.copy_(
                 torch.nn.functional.pad(
                     torch.cumsum(
@@ -956,6 +1057,30 @@ class FlashAttentionBackend(AttentionBackend):
             page_table = self.req_to_token[req_pool_indices, : metadata.max_seq_len_k]
             metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
 
+        if encoder_lens is not None:
+            # Only support encoder size 1 for now
+            metadata.encoder_max_seq_len_k = encoder_lens[0]
+            metadata.encoder_lens_int32.copy_(encoder_lens[:1])
+            metadata.encoder_cu_seqlens_k.copy_(
+                torch.nn.functional.pad(
+                    torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
+                    (1, 0),
+                )
+            )
+
+            metadata.encoder_page_table[:, : metadata.encoder_max_seq_len_k].copy_(
+                self.req_to_token[req_pool_indices, : metadata.encoder_max_seq_len_k]
+            )
+
+            # Update the regular page table
+            page_table = self.req_to_token[
+                req_pool_indices,
+                metadata.encoder_max_seq_len_k : (
+                    metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
+                ),
+            ]
+            metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+
         self.forward_metadata = metadata
 
     def get_cuda_graph_seq_len_fill_value(self):
@@ -972,14 +1097,19 @@ class FlashAttentionMultiStepBackend:
         self.topk = topk
         self.speculative_num_steps = speculative_num_steps
 
+        # TODO: Support Topk > 1 for FlashAttentionBackend Spec Decoding
+        assert (
+            self.topk == 1
+        ), "speculative_eagle_topk must be 1 for FlashAttentionMultiStepBackend"
+
         self.attn_backends = []
         for i in range(self.speculative_num_steps):
             self.attn_backends.append(
                 FlashAttentionBackend(
                     model_runner,
+                    speculative_step_id=i,
                     topk=self.topk,
                     speculative_num_steps=self.speculative_num_steps,
-                    step_id=i,
                 )
             )
 
@@ -1004,7 +1134,7 @@ class FlashAttentionMultiStepBackend:
                 forward_batch.batch_size * self.topk,
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
-                encoder_lens=None,
+                encoder_lens=forward_batch.encoder_lens,
                 forward_mode=ForwardMode.DECODE,
                 spec_info=forward_batch.spec_info,
             )
@@ -1021,7 +1151,7 @@ class FlashAttentionMultiStepBackend:
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
                 forward_batch.seq_lens_sum,
-                encoder_lens=None,
+                encoder_lens=forward_batch.encoder_lens,
                 forward_mode=ForwardMode.DECODE,
                 spec_info=forward_batch.spec_info,
                 seq_lens_cpu=forward_batch.seq_lens_cpu,