Dequant in chunked prefill

Signed-off-by: Matthew Bonanni <[email protected]>

Author: MatthewBonanni

vllm/v1/attention/backends/mla/common.py

@@ -427,6 +427,9 @@ def __init__(self,
             self.page_size = self.kv_cache_spec.block_size
 
         if self.chunked_prefill_enabled:
+            workspace_dtype = self.model_config.dtype
+            if cache_config.cache_dtype.startswith("fp8"):
+                workspace_dtype = current_platform.fp8_dtype()
             self.chunked_prefill_workspace_size = min(
                 # Max sure there is enough for 8 full length request or at least
                 # 4 pages of cache per request
@@ -447,7 +450,7 @@ def __init__(self,
             self.chunked_prefill_workspace = torch.empty(
                 (self.chunked_prefill_workspace_size,
                  self.model_config.get_head_size()),
-                dtype=self.model_config.dtype,
+                dtype=workspace_dtype,
                 device=device,
             )
 
@@ -1022,6 +1025,8 @@ def _compute_prefill_context(
         iters = len(prefill_metadata.chunked_context.seq_tot)
         workspace = prefill_metadata.chunked_context.workspace
 
+        fp8_attention = self.kv_cache_dtype.startswith("fp8")
+
         for i in range(iters):
             toks = prefill_metadata.chunked_context.seq_tot[i]
 
@@ -1039,6 +1044,16 @@ def _compute_prefill_context(
             k_pe = workspace[:toks]\
                 [..., self.kv_lora_rank:].unsqueeze(1)
 
+            if fp8_attention:
+                target_dtype = self.kv_b_proj.weight.dtype
+                kv_c_normed_dequant = torch.empty_like(kv_c_normed,
+                                                       dtype=target_dtype)
+                k_pe_dequant = torch.empty_like(k_pe, dtype=target_dtype)
+                ops.convert_fp8(kv_c_normed_dequant, kv_c_normed)
+                ops.convert_fp8(k_pe_dequant, k_pe)
+                kv_c_normed = kv_c_normed_dequant
+                k_pe = k_pe_dequant
+
             kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \
                 -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
             k_nope, v = kv_nope\
@@ -1155,7 +1170,7 @@ def forward(
             # same expert outputs.
             return output.fill_(0)
 
-        fp8_attention = self.kv_cache_dtype == "fp8"
+        fp8_attention = self.kv_cache_dtype.startswith("fp8")
 
         num_actual_toks = attn_metadata.num_actual_tokens
 
@@ -1191,6 +1206,9 @@ def forward(
                 scale=layer._k_scale,
             )
 
+        if fp8_attention:
+            kv_cache = kv_cache.view(current_platform.fp8_dtype())
+
         if has_prefill:
             output[num_decode_tokens:] = self._forward_prefill(
                 prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
@@ -1208,7 +1226,6 @@ def forward(
             decode_ql_nope = decode_ql_nope.transpose(0, 1)
 
             if fp8_attention:
-                kv_cache = kv_cache.view(torch.float8_e4m3fn)
                 ql_nope_shape = decode_ql_nope.shape
                 decode_ql_nope, _ = ops.scaled_fp8_quant(
                     decode_ql_nope.reshape([