sglang 0.4.6.post4__py3-none-any.whl → 0.4.6.post5__py3-none-any.whl

sglang/srt/layers/attention/aiter_backend.py

@@ -0,0 +1,513 @@
+from __future__ import annotations
+
+"""
+end to end attention solution with aiter kernels
+"""
+
+import math
+import os
+from dataclasses import dataclass
+from enum import Enum, auto
+from functools import partial
+from typing import TYPE_CHECKING, List, Optional, Union
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.speculative.spec_info import SpecInfo
+
+try:
+    from aiter import mha_batch_prefill_func, paged_attention_ragged
+except ImportError:
+    print(
+        "aiter is AMD specific kernel library. Please make sure aiter is installed on your AMD device."
+    )
+
+
+class WrapperDispatch(Enum):
+    SLIDING_WINDOW = auto()
+    CROSS_ATTENTION = auto()
+
+
+@dataclass
+class ForwardMetadata:
+    kv_indptr: torch.Tensor
+    kv_indices: torch.Tensor
+    max_q_len: int
+    max_kv_len: int
+
+
+global_workspace_buffer = None
+
+_AITER_PARTITION_SIZE_ROCM = 256
+
+
+class AiterAttnBackend(AttentionBackend):
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        kv_indptr_buf: Optional[torch.Tensor] = None,
+    ):
+        super().__init__()
+
+        self.device = model_runner.device
+        self.is_multimodal = model_runner.model_config.is_multimodal
+        self.num_head = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.head_dim = model_runner.model_config.head_dim
+        self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
+        self.num_kv_head = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.kv_cache_dtype = model_runner.kv_cache_dtype
+
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+
+        # Parse constants
+        self.max_context_len = model_runner.model_config.context_len
+        self.skip_prefill = skip_prefill
+
+        max_bs = model_runner.req_to_token_pool.size
+
+        if kv_indptr_buf is None:
+            self.kv_indptr = torch.zeros(
+                (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            self.kv_indptr = kv_indptr_buf
+
+        self.kv_last_page_len = torch.ones(
+            (max_bs,), dtype=torch.int32, device=model_runner.device
+        )
+        self.qo_indptr = torch.zeros(
+            (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+        )
+
+        # Create prefill indices updater
+        if not skip_prefill:
+            self.indices_updater_prefill = AiterIndicesUpdaterPrefill(
+                model_runner, self
+            )
+
+        # aiter kernel related initialization
+        self.max_num_partitions = (
+            self.max_context_len + _AITER_PARTITION_SIZE_ROCM - 1
+        ) // _AITER_PARTITION_SIZE_ROCM
+
+        nbyes_per_qo_elem = torch.finfo(torch.float32).bits // 8
+
+        self.workspace_buffer = torch.empty(
+            (max_bs * self.num_head * self.max_num_partitions * self.head_dim)
+            * nbyes_per_qo_elem
+            + 2 * (max_bs * self.num_head * self.max_num_partitions) * 4,
+            dtype=torch.uint8,
+            device=self.device,
+        )
+
+        self.scale = float(1.0 / (self.head_dim**0.5))
+        self.k_scale = self.v_scale = torch.tensor([1.0], dtype=torch.float32).to(
+            self.device
+        )
+        self.kv_last_page_lens = torch.ones((max_bs,), dtype=torch.int32).to(
+            self.device
+        )
+
+        self.logits_soft_cap = 0.0
+
+        self.forward_metadata: ForwardMetadata = None
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_decode_or_idle():
+            # update for aiter
+            # create kv_indices and kv_inptr
+            bs = forward_batch.batch_size
+            kv_indptr = self.kv_indptr
+            spec_info = forward_batch.spec_info
+            if spec_info is None:
+                kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                kv_indices = torch.zeros(
+                    forward_batch.seq_lens_sum, dtype=torch.int32, device=self.device
+                )
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    forward_batch.req_pool_indices,
+                    forward_batch.seq_lens,
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+            else:
+                kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+                bs = kv_indptr.shape[0] - 1
+
+            self.forward_metadata = ForwardMetadata(kv_indptr, kv_indices, None, None)
+
+        elif forward_batch.forward_mode.is_draft_extend():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = ForwardMetadata(
+                self.indices_updater_prefill.kv_indptr,
+                self.indices_updater_prefill.kv_indices,
+                self.indices_updater_prefill.max_q_len,
+                self.indices_updater_prefill.max_kv_len,
+            )
+        elif forward_batch.forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens=None,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=forward_batch.spec_info,
+            )
+            self.forward_metadata = ForwardMetadata(
+                self.indices_updater_prefill.kv_indptr,
+                self.indices_updater_prefill.kv_indices,
+                self.indices_updater_prefill.max_q_len,
+                self.indices_updater_prefill.max_kv_len,
+            )
+        else:
+            prefix_lens = forward_batch.extend_prefix_lens
+
+            if self.is_multimodal:
+                extend_no_prefix = False
+            else:
+                extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
+
+            self.indices_updater_prefill.update(
+                forward_batch.req_pool_indices,
+                forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
+                prefix_lens,
+                encoder_lens=forward_batch.encoder_lens,
+                spec_info=None,
+            )
+            self.forward_metadata = ForwardMetadata(
+                self.indices_updater_prefill.kv_indptr,
+                self.indices_updater_prefill.kv_indices,
+                self.indices_updater_prefill.max_q_len,
+                self.indices_updater_prefill.max_kv_len,
+            )
+
+    def init_cuda_graph_state(
+        self, max_bs: int, kv_indices_buf: Optional[torch.Tensor] = None
+    ):
+        if kv_indices_buf is None:
+            self.cuda_graph_kv_indices = torch.zeros(
+                (max_bs * self.max_context_len),
+                dtype=torch.int32,
+                device=self.device,
+            )
+        else:
+            self.cuda_graph_kv_indices = kv_indices_buf
+
+        if not self.skip_prefill:
+            self.cuda_graph_custom_mask = torch.zeros(
+                (max_bs * self.max_context_len),
+                dtype=torch.uint8,
+                device=self.device,
+            )
+
+    def init_forward_metadata_capture_cuda_graph(
+        self,
+        bs: int,
+        num_tokens: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+    ):
+        if forward_mode.is_decode_or_idle():
+            if spec_info is None:
+                kv_indptr = self.kv_indptr
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                kv_indices = self.cuda_graph_kv_indices
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices,
+                    seq_lens,
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+            else:
+                kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
+            self.forward_metadata = ForwardMetadata(kv_indptr, kv_indices, None, None)
+
+        elif forward_mode.is_target_verify():
+            seq_lens_sum = seq_lens.sum().item()
+            self.indices_updater_prefill.update(
+                req_pool_indices,
+                seq_lens,
+                seq_lens_sum,
+                prefix_lens=None,
+                encoder_lens=encoder_lens,
+                spec_info=spec_info,
+            )
+            self.forward_metadata = ForwardMetadata(
+                self.indices_updater_prefill.kv_indptr,
+                self.indices_updater_prefill.kv_indices,
+                self.indices_updater_prefill.max_q_len,
+                self.indices_updater_prefill.max_kv_len,
+            )
+
+        else:
+            raise ValueError(f"Invalid mode: {forward_mode=}")
+
+    def init_forward_metadata_replay_cuda_graph(
+        self,
+        bs: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        encoder_lens: Optional[torch.Tensor],
+        forward_mode: ForwardMode,
+        spec_info: Optional[SpecInfo],
+        seq_lens_cpu: Optional[torch.Tensor],
+    ):
+        if forward_mode.is_decode_or_idle():
+            kv_indptr = self.kv_indptr
+            kv_indices = self.cuda_graph_kv_indices
+            if spec_info is None:
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0)
+                kv_indptr = kv_indptr[: bs + 1]
+                create_flashinfer_kv_indices_triton[(bs,)](
+                    self.req_to_token,
+                    req_pool_indices[:bs],
+                    seq_lens[:bs],
+                    kv_indptr,
+                    None,
+                    kv_indices,
+                    self.req_to_token.stride(0),
+                )
+            else:
+                kv_indptr[: spec_info.kv_indptr.shape[0]] = spec_info.kv_indptr
+                kv_indices[: spec_info.kv_indices.shape[0]] = spec_info.kv_indices
+
+        elif forward_mode.is_target_verify():
+            self.indices_updater_prefill.update(
+                req_pool_indices[:bs],
+                seq_lens[:bs],
+                seq_lens_sum,
+                prefix_lens=None,
+                encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
+                spec_info=spec_info,
+            )
+        else:
+            raise ValueError("Invalid forward mode")
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 1
+
+    def forward_extend(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
+
+        self.logits_soft_cap = layer.logit_cap
+
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, cache_loc, k, v, layer.k_scale, layer.v_scale
+                )
+
+        k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
+
+        bs0 = forward_batch.batch_size + 1
+
+        o = mha_batch_prefill_func(
+            q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+            k_cache,
+            v_cache,
+            self.qo_indptr[:bs0],
+            self.forward_metadata.kv_indptr[:bs0],
+            self.forward_metadata.kv_indices,
+            self.forward_metadata.max_q_len,
+            self.forward_metadata.max_kv_len,
+            causal=True,
+            logits_soft_cap=self.logits_soft_cap,
+            alibi_slopes=None,
+            return_lse=False,
+            return_attn_probs=False,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+
+    def forward_decode(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
+    ):
+        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
+
+        if layer.qk_head_dim != layer.v_head_dim:
+            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
+
+        self.logits_soft_cap = layer.logit_cap
+        paged_attention_ragged(
+            o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            self.workspace_buffer,
+            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
+            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).view(
+                -1, 1, layer.tp_k_head_num, layer.qk_head_dim
+            ),
+            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id).view(
+                -1, 1, layer.tp_v_head_num, layer.v_head_dim
+            ),
+            self.scale,
+            self.forward_metadata.kv_indptr,
+            self.forward_metadata.kv_indices,
+            self.kv_last_page_lens,
+            1,
+            self.max_num_partitions,
+            None,
+            "auto",
+            "NHD",
+            self.logits_soft_cap,
+            self.k_scale,
+            self.v_scale,
+            None,
+            _AITER_PARTITION_SIZE_ROCM,
+        )
+
+        return o
+
+
+class AiterIndicesUpdaterPrefill:
+    def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
+        # Parse Constants
+        self.num_qo_heads = (
+            model_runner.model_config.num_attention_heads // get_attention_tp_size()
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            get_attention_tp_size()
+        )
+        self.head_dim = model_runner.model_config.head_dim
+        self.data_type = model_runner.kv_cache_dtype
+        self.q_data_type = model_runner.dtype
+        self.sliding_window_size = model_runner.sliding_window_size
+        self.attn_backend = attn_backend
+
+        # Buffers and wrappers
+        self.kv_indptr = attn_backend.kv_indptr
+        self.kv_last_page_len = attn_backend.kv_last_page_len
+        self.qo_indptr = attn_backend.qo_indptr
+        self.req_to_token = model_runner.req_to_token_pool.req_to_token
+        self.update = self.update_single_wrapper
+
+        self.kv_indices = None
+        self.max_q_len = 0
+        self.max_kv_len = 0
+
+    def update(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[SpecInfo],
+    ):
+        # Keep the signature for type checking. It will be assigned during runtime.
+        raise NotImplementedError()
+
+    def update_single_wrapper(
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor],
+        spec_info: Optional[SpecInfo],
+    ):
+
+        kv_start_idx = None
+        kv_indptr = self.kv_indptr
+        qo_indptr = self.qo_indptr
+        paged_kernel_lens = seq_lens
+        paged_kernel_lens_sum = seq_lens_sum
+
+        bs = len(req_pool_indices)
+        if spec_info is None:
+            # Normal extend
+            kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
+            kv_indptr = kv_indptr[: bs + 1]
+            kv_indices = torch.empty(
+                paged_kernel_lens_sum + 256,
+                dtype=torch.int32,
+                device=req_pool_indices.device,
+            )
+            create_flashinfer_kv_indices_triton[(bs,)](
+                self.req_to_token,
+                req_pool_indices,
+                paged_kernel_lens,
+                kv_indptr,
+                kv_start_idx,
+                kv_indices,
+                self.req_to_token.shape[1],
+            )
+
+            self.max_kv_len = torch.max(paged_kernel_lens).item()
+
+            extend_lens = seq_lens - prefix_lens
+            self.max_q_len = torch.max(extend_lens).item()
+
+            qo_indptr[1 : bs + 1] = torch.cumsum(extend_lens, dim=0)
+            qo_indptr = qo_indptr[: bs + 1]
+            custom_mask = None
+        else:
+            kv_indices, kv_indptr, qo_indptr, custom_mask = (
+                spec_info.generate_attn_arg_prefill(
+                    req_pool_indices,
+                    paged_kernel_lens,
+                    self.req_to_token,
+                )
+            )
+
+        self.kv_indices = kv_indices

sglang/srt/layers/attention/flashattention_backend.py

@@ -918,8 +918,11 @@ class FlashAttentionBackend(AttentionBackend):
             and local_attn_metadata is not None
             and (hasattr(layer, "use_irope") and layer.use_irope)
         )
-        # We do cascade attention for Draft Decode with topk > 1
-        use_cascade_attn = self.topk > 1
+
+        # When Spec Decode enabled, forward_decode would be called with two mode:
+        # 1. DRAFT_DECODE: we enable cascade attention when top_k > 1
+        # 2. IDLE: we don’t need cascade attention, spec_info will be none in this case
+        use_cascade_attn = forward_batch.spec_info is not None and self.topk > 1
 
         # Calculate window size (can be moved to metadata if layer properties don't change)
         # we don't do layer.sliding_window_size - 1 since in model.get_attention_sliding_window_size() we already - 1
@@ -1165,7 +1168,6 @@ class FlashAttentionBackend(AttentionBackend):
             max_virtual_batches = max_bs * (
                 (max_seq_len + attn_chunk_size - 1) // attn_chunk_size
             )
-            max_blocks_per_seq = (max_seq_len + attn_chunk_size - 1) // attn_chunk_size
             max_pages_per_block = (attn_chunk_size + page_size - 1) // page_size
 
             self.decode_cuda_graph_local_attn_metadata = {
@@ -1177,7 +1179,7 @@ class FlashAttentionBackend(AttentionBackend):
                 ),
                 "local_block_table": torch.zeros(
                     max_virtual_batches,
-                    max_blocks_per_seq * max_pages_per_block,
+                    max_pages_per_block,
                     dtype=torch.int32,
                     device=self.device,
                 ),
@@ -1435,19 +1437,7 @@ class FlashAttentionBackend(AttentionBackend):
                 self.decode_cuda_graph_metadata[bs] = metadata
 
                 if self.attention_chunk_size is not None:
-                    metadata.local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
-                        local_query_start_loc=self.decode_cuda_graph_local_attn_metadata[
-                            "local_query_start_loc"
-                        ],
-                        local_seqused_k=self.decode_cuda_graph_local_attn_metadata[
-                            "local_seqused_k"
-                        ],
-                        local_block_table=self.decode_cuda_graph_local_attn_metadata[
-                            "local_block_table"
-                        ],
-                        local_max_query_len=1,
-                        local_max_seq_len=1,
-                    )
+                    self._update_local_attn_metadata_for_capture(metadata, batch_size)
 
         elif forward_mode.is_target_verify():
             if self.topk <= 1:
@@ -1808,6 +1798,62 @@ class FlashAttentionBackend(AttentionBackend):
         )
         metadata.local_attn_metadata = local_metadata
 
+    def _update_local_attn_metadata_for_capture(
+        self, metadata: FlashAttentionMetadata, bs: int
+    ):
+        """Update local attention metadata during CUDA graph capture phase.
+
+        This method calculates the exact buffer sizes needed for local attention metadata
+        during the CUDA graph capture phase, optimizing memory usage by creating views of
+        pre-allocated buffers with exactly the sizes needed.
+        """
+        seq_lens_capture = metadata.cache_seqlens_int32
+        max_seq_len = int(seq_lens_capture.max().item())
+        page_table_capture = metadata.page_table
+
+        cu_seqlens_q_np = metadata.cu_seqlens_q.cpu().numpy()
+        seqlens_np = seq_lens_capture.cpu().numpy()
+        (
+            seqlens_q_local_np,
+            cu_seqlens_q_local_np,
+            seqlens_k_local_np,
+            block_table_local_np,
+        ) = make_local_attention_virtual_batches(
+            self.attention_chunk_size,
+            cu_seqlens_q_np,
+            seqlens_np,
+            page_table_capture,
+            self.page_size,
+        )
+
+        # Get exact dimensions from the calculation
+        q_len = len(cu_seqlens_q_local_np)
+        k_len = len(seqlens_k_local_np)
+        b0 = block_table_local_np.shape[0] if block_table_local_np.shape[0] > 0 else bs
+        b1 = block_table_local_np.shape[1] if block_table_local_np.shape[1] > 0 else 1
+
+        # Create views of the pre-allocated buffers with exactly these sizes
+        # This is the key optimization - we only use the memory we actually need
+        local_query_start_loc = self.decode_cuda_graph_local_attn_metadata[
+            "local_query_start_loc"
+        ][:q_len]
+
+        local_seqused_k = self.decode_cuda_graph_local_attn_metadata["local_seqused_k"][
+            :k_len
+        ]
+
+        local_block_table = self.decode_cuda_graph_local_attn_metadata[
+            "local_block_table"
+        ][:b0, :b1]
+
+        metadata.local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
+            local_query_start_loc=local_query_start_loc,
+            local_seqused_k=local_seqused_k,
+            local_block_table=local_block_table,
+            local_max_query_len=1,
+            local_max_seq_len=max_seq_len,
+        )
+
     def _update_local_attn_metadata_for_replay(
         self, metadata: FlashAttentionMetadata, bs: int
     ):

sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -346,7 +346,6 @@ class FlashInferMLAAttnBackend(AttentionBackend):
         cache_loc = forward_batch.out_cache_loc
         logits_soft_cap = layer.logit_cap
         prefill_wrapper_paged = self.forward_metadata.prefill_wrapper
-        k_buf = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
 
         # Save kv cache
         if save_kv_cache and k is not None:
@@ -381,6 +380,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             )
         else:
             # mla paged prefill
+            k_buf = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to(
+                q.dtype
+            )
             if q_rope is None:
                 qall = q.view(-1, layer.tp_q_head_num, layer.head_dim)
                 q, q_rope = (
@@ -442,7 +444,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
             q_nope = reshaped_q[:, :, : layer.v_head_dim]
             q_rope = reshaped_q[:, :, layer.v_head_dim :]
 
-        k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
+        k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id).to(
+            q.dtype
+        )
 
         o = q_nope.new_empty(q_nope.shape)
         # Direct call to run without the wrapper
@@ -467,7 +471,7 @@ class FlashInferMLAIndicesUpdaterDecode:
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
         self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
         self.scaling = model_runner.model_config.scaling
-        self.data_type = model_runner.kv_cache_dtype
+        self.data_type = model_runner.dtype
         self.attn_backend = attn_backend
 
         # Buffers and wrappers
@@ -577,7 +581,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
         self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
         self.v_head_dim = model_runner.model_config.v_head_dim
         self.scaling = model_runner.model_config.scaling
-        self.data_type = model_runner.kv_cache_dtype
+        self.data_type = model_runner.dtype
         self.q_data_type = model_runner.dtype
         self.attn_backend = attn_backend