[Intel GPU] DeepSeek V4 6/N: use sgl-kernel implemetation of flash_mla_with_kvcache on XPU (#27914)
Signed-off-by: P V R K Jyothendra Varma <polisetty.v.r.k.jyothendra.varma@intel.com> Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
co-authored by
Ma Mingfei
parent
e90fec4868
commit
4dddb04325
@@ -50,7 +50,7 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
|||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||||
from sglang.srt.runtime_context import get_parallel
|
from sglang.srt.runtime_context import get_parallel
|
||||||
from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
|
from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
|
||||||
from sglang.srt.utils import ceil_align
|
from sglang.srt.utils import ceil_align, is_xpu
|
||||||
from sglang.srt.utils.common import is_sm120_supported
|
from sglang.srt.utils.common import is_sm120_supported
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -60,6 +60,7 @@ if TYPE_CHECKING:
|
|||||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||||
|
|
||||||
_is_sm120 = is_sm120_supported()
|
_is_sm120 = is_sm120_supported()
|
||||||
|
_is_xpu = is_xpu()
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -111,7 +112,7 @@ def _pad_last_dim(x: T, multiples_of: int = PAGE_INDEX_ALIGNED_SIZE) -> T:
|
|||||||
|
|
||||||
|
|
||||||
def _create_flashmla_metadata():
|
def _create_flashmla_metadata():
|
||||||
if _is_sm120:
|
if _is_sm120 or _is_xpu:
|
||||||
return None
|
return None
|
||||||
import sgl_kernel.flash_mla as flash_mla
|
import sgl_kernel.flash_mla as flash_mla
|
||||||
|
|
||||||
@@ -1433,9 +1434,12 @@ class DeepseekV4AttnBackend(
|
|||||||
extra_topk_length=extra_topk_lengths,
|
extra_topk_length=extra_topk_lengths,
|
||||||
)[0]
|
)[0]
|
||||||
else:
|
else:
|
||||||
import sgl_kernel.flash_mla as flash_mla
|
if _is_xpu:
|
||||||
|
from sgl_kernel import flash_mla_with_kvcache
|
||||||
|
else:
|
||||||
|
from sgl_kernel.flash_mla import flash_mla_with_kvcache
|
||||||
|
|
||||||
o = flash_mla.flash_mla_with_kvcache(
|
o = flash_mla_with_kvcache(
|
||||||
q=q,
|
q=q,
|
||||||
k_cache=swa_k_cache,
|
k_cache=swa_k_cache,
|
||||||
head_dim_v=self.head_dim_v,
|
head_dim_v=self.head_dim_v,
|
||||||
|
|||||||
Reference in New Issue
Block a user