[XPU] Fix incorrect context-key normalization for Qwen DFlash-based models (#56431)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Sep 18, 2026 at 07:52 UTC
48f663c37f8c77803d9281778d49ec0727abcacf
1 file changed
+10
vllm/model_executor/models/qwen3_dflash.py
+10
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
32
VocabParallelEmbedding,
33
)
34
from vllm.multimodal.inputs import NestedTensors
35
+from vllm.platforms import current_platform
36
from vllm.transformers_utils.config import set_default_rope_theta
37
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
38
from vllm.v1.attention.backend import AttentionType
@@ -561,6 +562,15 @@ class DFlashQwen3Model(nn.Module):
562
# --- Grouped RMSNorm K across all layers ([L, num_ctx, nkv, hd]) ---
563
# The weight is selected per layer by the outermost (layer) index.
564
all_k_normed = torch.empty_like(all_k)
565
+ if current_platform.is_xpu():
566
+ for layer_idx in range(all_k.shape[0]):
567
+ ops.rms_norm(
568
+ all_k_normed[layer_idx],
569
+ all_k[layer_idx],
570
+ self._k_norm_weights[layer_idx],
571
+ self._rms_norm_eps,
572
+ )
573
+ return all_k_normed
574
ops.rms_norm(
575
all_k_normed,
576
all_k,