[XPU] use rms_norm xpu kernel for context-key normalization (#58896)
Signed-off-by: Yan Ma <yan.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yan Ma committed
Sep 29, 2026 at 07:09 UTC
754aa29c006ddfbc518ae7d89e30414c6866d890
1 file changed
-10
vllm/model_executor/models/qwen3_dflash.py
-10
@@ -34,7 +34,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
34
VocabParallelEmbedding,
35
)
36
from vllm.multimodal.inputs import NestedTensors
37
-from vllm.platforms import current_platform
37
from vllm.transformers_utils.config import set_default_rope_theta
38
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
39
from vllm.v1.attention.backend import AttentionType
@@ -622,15 +621,6 @@ class DFlashQwen3Model(nn.Module):
621
# --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
622
# The weight is selected per layer by the outermost (layer) index.
623
all_k_normed = torch.empty_like(all_k)
625
- if current_platform.is_xpu():
626
- for layer_idx in range(all_k.shape[0]):
627
- ops.rms_norm(
628
- all_k_normed[layer_idx],
629
- all_k[layer_idx],
630
- self._k_norm_weights[layer_idx],
631
- self._rms_norm_eps,
632
- )
633
- return all_k_normed
624
ops.rms_norm(
625
all_k_normed,
626
all_k,