[Feat] Support FlashMLA backend with MTP and FP8 KV cache

python/sglang/srt/layers/attention/flashinfer_mla_backend.py

@@ -346,7 +346,6 @@ def forward_extend(
         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 @@ def forward_extend(
             )
         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 @@ def forward_decode(
             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 @@ def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         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 @@ def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
         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