support fuse cat + q to fp8 + mla

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

@@ -1318,7 +1318,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
                             False,
                             1e-6,
                         ) 
-                    else:
+                    if has_decode:
                         q_tensor = torch.randn(q.shape[0], num_local_heads, self.qk_nope_head_dim + self.qk_rope_head_dim, dtype=q.dtype, device=q.device)
                         q_quant = torch.empty_like(q_tensor, dtype=torch.float8_e4m3fn, device=q.device)
                         q_scale = torch.empty(q.shape[0], dtype=torch.float32, device=q.device)

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

@@ -186,7 +186,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
                 o, _ = flash_mla_with_kvcache_fp8_with_cat(
                     q_nope=q_nope.unsqueeze(1),
                     q_pe=q_pe.unsqueeze(1),
-                    k_cache=kv_c_and_k_pe_cache.unsqueeze(-2),  # Add head dim of 1
+                    k_cache=kv_c_and_k_pe_cache.unsqueeze(-2).view(torch.float8_e4m3fn),  # Add head dim of 1
                     block_table=attn_metadata.decode.block_table,
                     cache_seqlens=attn_metadata.decode.seq_lens,
                     head_dim_v=self.kv_lora_rank,
@@ -199,32 +199,49 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
                     descale_k = k_scale,            
                 )
             else:
-                if envs.VLLM_USE_OPT_CAT:
-                    if q_nope.shape[0] < 1024:
-                        from vllm.v1.attention.backends.mla.test_concat import concat_helper_decode
-                        q = concat_helper_decode(q_nope, q_pe, dim=2)\
-                            .unsqueeze(1)
+                if envs.VLLM_USE_CAT_MLA:
+                    o, _ = flash_mla_with_kvcache_fp8_with_cat(
+                        q_nope=q_nope.unsqueeze(1),
+                        q_pe=q_pe.unsqueeze(1),
+                        k_cache=kv_c_and_k_pe_cache.unsqueeze(-2).view(torch.float8_e4m3fn),  # Add head dim of 1
+                        block_table=attn_metadata.decode.block_table,
+                        cache_seqlens=attn_metadata.decode.seq_lens,
+                        head_dim_v=self.kv_lora_rank,
+                        tile_scheduler_metadata=attn_metadata.decode.
+                        tile_scheduler_metadata,
+                        num_splits=attn_metadata.decode.num_splits,
+                        softmax_scale=self.scale,
+                        causal=True,
+                        descale_q = q_scale,
+                        descale_k = k_scale,            
+                    )
+                else:
+                    if envs.VLLM_USE_OPT_CAT:
+                        if q_nope.shape[0] < 1024:
+                            from vllm.v1.attention.backends.mla.test_concat import concat_helper_decode
+                            q = concat_helper_decode(q_nope, q_pe, dim=2)\
+                                .unsqueeze(1)
+                        else:
+                            q = torch.cat([q_nope, q_pe], dim=-1)\
+                            .unsqueeze(1) # Add seqlen dim of 1 (decode)
                     else:
                         q = torch.cat([q_nope, q_pe], dim=-1)\
-                        .unsqueeze(1) # Add seqlen dim of 1 (decode)
-                else:
-                    q = torch.cat([q_nope, q_pe], dim=-1)\
-                        .unsqueeze(1) # Add seqlen dim of 1 (decode)
-                o, _ = flash_mla_with_kvcache_fp8(
-                    q=q.to(torch.float8_e4m3fn),
-                    k_cache=kv_c_and_k_pe_cache.unsqueeze(-2).view(torch.float8_e4m3fn),  # Add head dim of 1
-                    block_table=attn_metadata.decode.block_table,
-                    cache_seqlens=attn_metadata.decode.seq_lens,
-                    head_dim_v=self.kv_lora_rank,
-                    tile_scheduler_metadata=attn_metadata.decode.
-                    tile_scheduler_metadata,
-                    num_splits=attn_metadata.decode.num_splits,
-                    softmax_scale=self.scale,
-                    causal=True,
-                    descale_q=q_scale,
-                    descale_k=k_scale,         
-                )
-            
+                            .unsqueeze(1) # Add seqlen dim of 1 (decode)
+                    o, _ = flash_mla_with_kvcache_fp8(
+                        q=q.to(torch.float8_e4m3fn),
+                        k_cache=kv_c_and_k_pe_cache.unsqueeze(-2).view(torch.float8_e4m3fn),  # Add head dim of 1
+                        block_table=attn_metadata.decode.block_table,
+                        cache_seqlens=attn_metadata.decode.seq_lens,
+                        head_dim_v=self.kv_lora_rank,
+                        tile_scheduler_metadata=attn_metadata.decode.
+                        tile_scheduler_metadata,
+                        num_splits=attn_metadata.decode.num_splits,
+                        softmax_scale=self.scale,
+                        causal=True,
+                        descale_q=q_scale,
+                        descale_k=k_scale,         
+                    )
+                    
         else:
             if not envs.VLLM_USE_CAT_MLA or kv_cache_dtype == "fp8_e4m3":
                 if envs.VLLM_USE_OPT_CAT: