minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Intel GPU] DeepSeek V4 8/N: use sgl-kernel implementation of fused_k_norm_rope_flashmla on XPU (#28040)

Browse Source 
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:
Polisetty V R K Jyothendra Varma
2026-08-05 09:28:43 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 9303e26f03
commit d2c405f19d
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsv4/elementwise.py
+8
View File
@@ -15,6 +15,9 @@ from .utils import make_name
_is_hip = is_hip() _is_hip = is_hip()
_is_xpu = is_xpu() _is_xpu = is_xpu()
if _is_xpu:
from sgl_kernel import fused_k_norm_rope_flashmla as fused_k_norm_rope_flashmla_xpu
@cache_once @cache_once
def _jit_fused_rope_module(): def _jit_fused_rope_module():
@@ -270,6 +273,11 @@ def fused_k_norm_rope_flashmla(
freqs_real = torch.view_as_real(freqs_cis).flatten(-2) freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
head_dim = kv.shape[-1] head_dim = kv.shape[-1]
rope_dim = freqs_real.shape[-1] rope_dim = freqs_real.shape[-1]
if _is_xpu:
fused_k_norm_rope_flashmla_xpu(
kv, kv_weight, freqs_real, positions, out_loc, kvcache, eps, page_size
)
else:
module = _jit_main_k_norm_rope_flashmla_module( module = _jit_main_k_norm_rope_flashmla_module(
kv.dtype, head_dim, rope_dim, page_size kv.dtype, head_dim, rope_dim, page_size
) )