Enable deepseek-r1 with xpu-kernels. (vllm-project#6)

* [kernel][DS-R1][linear] use default Fp8LinearMethod/Fp8MoEMethod

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

* [kernel][DS-R1][Attention] enable Triton MLA attention

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

* enable MHA for deepseek, need padding head_size to make flash attn kernel happy

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

* not break fp8 path

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

---------

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

Author: jikunshang

vllm/_ipex_ops.py

@@ -62,6 +62,7 @@ def flash_attn_varlen_func(
         k_descale=None,
         v_descale=None,
         num_splits=0,
+        return_softmax_lse: bool | None = False,
         s_aux: torch.Tensor | None = None,
     ):
         if out is None:
@@ -97,6 +98,7 @@ def flash_attn_varlen_func(
             window_size=real_window_size,
             # alibi_slopes = alibi_slopes,
             # softcap=softcap,
+            return_softmax_lse=return_softmax_lse,
         )
 
     @staticmethod

vllm/attention/ops/triton_flash_attention.py

@@ -341,7 +341,19 @@ def get_rdna_autotune_configs():
     ], ["IS_CAUSAL", "dropout_p", "BLOCK_DMODEL", "USE_FP8"]
 
 
+def get_xpu_autotune_configs():
+    return [
+        triton.Config(
+            {"BLOCK_M": 32, "BLOCK_N": 32, "PRE_LOAD_V": False},
+            num_stages=1,
+            num_warps=2,
+        ),
+    ], ["IS_CAUSAL", "dropout_p", "BLOCK_DMODEL", "USE_FP8"]
+
+
 def get_autotune_configs():
+    if current_platform.is_xpu():
+        return get_xpu_autotune_configs()
     if on_gfx1x():
         return get_rdna_autotune_configs()
     else:

vllm/envs.py

@@ -224,6 +224,7 @@
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_FLATTEN_LOGPROBS: bool = False
     VLLM_XPU_USE_W8A8_GEMM: bool = False
+    VLLM_XPU_ATTN_HEAD_SIZE_PAD: bool = False
 
 
 def get_default_cache_root():
@@ -1486,6 +1487,9 @@ def get_vllm_port() -> int | None:
     "VLLM_XPU_USE_W8A8_GEMM": lambda: bool(
         int(os.getenv("VLLM_XPU_USE_W8A8_GEMM", "0"))
     ),
+    "VLLM_XPU_ATTN_HEAD_SIZE_PAD": lambda: bool(
+        int(os.getenv("VLLM_XPU_ATTN_HEAD_SIZE_PAD", "0"))
+    ),
 }
 
 # --8<-- [end:env-vars-definition]
@@ -1613,6 +1617,7 @@ def compute_hash() -> str:
         "VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE",
         "VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL",
         "VLLM_XPU_USE_W8A8_GEMM",
+        "VLLM_XPU_ATTN_HEAD_SIZE_PAD",
     ]
     for key in environment_variables_to_hash:
         # if this goes out of sync with environment_variables,

vllm/model_executor/layers/quantization/fp8.py

@@ -153,7 +153,7 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:
         not current_platform.has_device_capability(89)
         or envs.VLLM_TEST_FORCE_FP8_MARLIN
     )
-    if current_platform.is_rocm():
+    if current_platform.is_rocm() or current_platform.is_xpu():
         use_marlin = False
     if use_marlin:
         logger.info_once("Using Marlin backend for FP8 MoE")
@@ -284,7 +284,9 @@ def get_quant_method(
     ) -> Optional["QuantizeMethodBase"]:
         from vllm.attention.layer import Attention  # Avoid circular import
 
-        if current_platform.is_xpu():
+        # for non-block quant on xpu, we use the xpu fp8 method,
+        # otherwise use triton
+        if current_platform.is_xpu() and self.weight_block_size is None:
             return self.get_xpu_quant_method(layer, prefix)
         if isinstance(layer, LinearBase):
             if is_layer_skipped(

vllm/platforms/xpu.py

@@ -73,6 +73,10 @@ def get_attn_backend_cls(
             raise NotImplementedError("Sparse Attention is not supported on XPU.")
         TRITON_ATTN = "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"  # noqa: E501
         FLASH_ATTN = "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend"  # noqa: E501
+        TRITON_ATTN_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"  # noqa: E501
+        if use_mla:
+            logger.info_once("Using Triton MLA backend on V1 engine.")
+            return TRITON_ATTN_MLA
         if selected_backend == _Backend.TRITON_ATTN:
             logger.info_once("Using Triton backend.")
             return TRITON_ATTN

vllm/v1/attention/backends/flash_attn.py

@@ -550,6 +550,31 @@ def forward(
             return output.fill_(0)
 
         attn_type = self.attn_type
+        output_may_pad = output  # default
+
+        if envs.VLLM_XPU_ATTN_HEAD_SIZE_PAD:
+            logger.warning_once(
+                "VLLM_XPU_ATTN_HEAD_SIZE_PAD is enabled. "
+                "Padding head size to 256 for FlashAttention."
+            )
+            # due to attention head size limitations in current flash attention
+            # kernel(which support 64/128/256 only), we will pad the head size
+            # to 256 for deepseek model.
+            orig_head_size = query.shape[-1]
+            new_shape = query.shape[:-1] + (256,)
+
+            query_pad = query.new_zeros(new_shape)
+            query_pad[..., : query.shape[-1]] = query
+            key_pad = key.new_zeros(new_shape)
+            key_pad[..., : key.shape[-1]] = key
+            value_pad = value.new_zeros(new_shape)
+            value_pad[..., : value.shape[-1]] = value
+            # for output, it's inplace?
+            output_may_pad = output.new_zeros(new_shape)
+
+            query = query_pad
+            key = key_pad
+            value = value_pad
 
         # IMPORTANT!
         # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
@@ -641,7 +666,7 @@ def forward(
                     q=query[:num_actual_tokens],
                     k=key_cache,
                     v=value_cache,
-                    out=output[:num_actual_tokens],
+                    out=output_may_pad[:num_actual_tokens],
                     cu_seqlens_q=cu_seqlens_q,
                     max_seqlen_q=max_seqlen_q,
                     seqused_k=seqused_k,
@@ -660,7 +685,12 @@ def forward(
                     num_splits=attn_metadata.max_num_splits,
                     s_aux=self.sinks,
                 )
-                return output
+                if envs.VLLM_XPU_ATTN_HEAD_SIZE_PAD:
+                    # it's inplace, we should not replace.
+                    output[:num_actual_tokens] = output_may_pad[
+                        :num_actual_tokens, :, :orig_head_size
+                    ]
+            return output
 
         # Cascade attention (rare case).
         cascade_attention(

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

@@ -251,7 +251,7 @@ class QueryLenSupport(Enum):
 
 
 try:
-    from vllm.vllm_flash_attn import flash_attn_varlen_func
+    from vllm.attention.utils.fa_utils import flash_attn_varlen_func
 
     is_vllm_fa = True
 except ImportError:

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

@@ -14,7 +14,9 @@
 )
 from vllm.attention.utils.fa_utils import (
     flash_attn_supports_mla,
+    flash_attn_varlen_func,
     get_flash_attn_version,
+    get_scheduler_metadata,
 )
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
@@ -31,7 +33,6 @@
 )
 from vllm.v1.attention.backends.utils import AttentionCGSupport
 from vllm.v1.kv_cache_interface import AttentionSpec
-from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
 
 logger = init_logger(__name__)
 

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

@@ -121,7 +121,7 @@ def _flash_attn_varlen_diff_headdims(
         self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs
     ):
         if (
-            current_platform.is_rocm()
+            (current_platform.is_rocm() or current_platform.is_xpu())
             and self.use_triton_flash_attn
             and not return_softmax_lse
         ):