A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 14:45 UTC
[XPU] Initial support for GDN attention on Qwen3-next/Qwen3.5 (#33657)
Signed-off-by: Yan Ma <yan.ma@intel.com> Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yan Ma committed
Apr 3, 2026 at 08:59 UTC
ee3cf457398e89f6c8cc2320c1cbdcfa362681cb
3 files changed
+150
vllm/model_executor/layers/layernorm.py
+5
@@ -560,6 +560,11 @@ class RMSNormGated(CustomOp):
560
activation=self.activation,
561
)
562
563
+ def forward_xpu(
564
+ self, x: torch.Tensor, z: torch.Tensor | None = None
565
+ ) -> torch.Tensor:
566
+ return self.forward_cuda(x, z)
567
+
568
569
class LayerNorm(nn.Module):
570
"""
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+94
@@ -262,6 +262,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
262
else 0
263
)
264
self.gqa_interleaved_layout = gqa_interleaved_layout
265
+ self._forward_method = (
266
+ self.forward_xpu if current_platform.is_xpu() else self.forward_cuda
267
+ )
268
269
# QKV
270
self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -493,6 +496,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
496
self,
497
hidden_states: torch.Tensor,
498
output: torch.Tensor,
499
+ ):
500
+ self._forward_method(hidden_states, output)
501
+
502
+ def forward_cuda(
503
+ self,
504
+ hidden_states: torch.Tensor,
505
+ output: torch.Tensor,
506
):
507
"""
508
Forward pass with three parts:
@@ -567,6 +577,90 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
577
core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
578
output[:num_tokens], _ = self.out_proj(core_attn_out)
579
580
+ def forward_xpu(
581
+ self,
582
+ hidden_states: torch.Tensor,
583
+ output: torch.Tensor,
584
+ ):
585
+ """
586
+ Forward pass with three parts:
587
+ 1. Input projection
588
+ 2. Core attention (custom op)
589
+ 3. Output projection
590
+ """
591
+ num_tokens = hidden_states.size(0)
592
+
593
+ assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on XPU."
594
+
595
+ # ============================================================
596
+ # Part 1: Input Projection
597
+ # ============================================================
598
+ projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
599
+ projected_states_ba, _ = self.in_proj_ba(hidden_states)
600
+
601
+ # ============================================================
602
+ # Part 2: Core Attention
603
+ # ============================================================
604
+ forward_context = get_forward_context()
605
+ attn_metadata: AttentionMetadata = forward_context.attn_metadata
606
+ core_attn_out = torch.zeros(
607
+ (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
608
+ dtype=hidden_states.dtype,
609
+ device=hidden_states.device,
610
+ )
611
+ z = torch.empty_like(core_attn_out)
612
+ if attn_metadata is not None:
613
+ attn_metadata = attn_metadata[self.prefix]
614
+
615
+ # TODO: xpu does not support this param yet
616
+ spec_sequence_masks = attn_metadata.spec_sequence_masks
617
+ assert spec_sequence_masks is None
618
+
619
+ conv_weights = self.conv1d.weight.view(
620
+ self.conv1d.weight.size(0), self.conv1d.weight.size(2)
621
+ )
622
+
623
+ conv_state = self.kv_cache[0]
624
+ ssm_state = self.kv_cache[1]
625
+
626
+ torch.ops._xpu_C.gdn_attention(
627
+ core_attn_out,
628
+ z,
629
+ projected_states_qkvz,
630
+ projected_states_ba,
631
+ self.num_k_heads,
632
+ self.num_v_heads,
633
+ self.head_k_dim,
634
+ self.head_v_dim,
635
+ conv_state=conv_state,
636
+ ssm_state=ssm_state,
637
+ conv_weights=conv_weights,
638
+ conv_bias=self.conv1d.bias,
639
+ activation=self.activation,
640
+ A_log=self.A_log,
641
+ dt_bias=self.dt_bias,
642
+ num_prefills=attn_metadata.num_prefills,
643
+ num_decodes=attn_metadata.num_decodes,
644
+ has_initial_state=attn_metadata.has_initial_state,
645
+ non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc,
646
+ non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor,
647
+ num_actual_tokens=attn_metadata.num_actual_tokens,
648
+ tp_size=self.tp_size,
649
+ reorder_input=not self.gqa_interleaved_layout,
650
+ )
651
+
652
+ # ============================================================
653
+ # Part 3: Output Projection
654
+ # ============================================================
655
+ z_shape_og = z.shape
656
+ # Reshape input data into 2D tensor
657
+ core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
658
+ z = z.reshape(-1, z.shape[-1])
659
+ core_attn_out = self.norm(core_attn_out, z)
660
+ core_attn_out = core_attn_out.reshape(z_shape_og)
661
+ core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
662
+ output[:num_tokens], _ = self.out_proj(core_attn_out)
663
+
664
def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
665
"""Warm up GDN prefill kernels during V1 profiling.
666
vllm/platforms/xpu.py
+51
@@ -218,6 +218,57 @@ class XPUPlatform(Platform):
218
# ref. https://openucx.readthedocs.io/en/master/faq.html
219
os.environ["UCX_MEMTYPE_CACHE"] = "n"
220
221
+ @classmethod
222
+ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
223
+ super().update_block_size_for_backend(vllm_config)
224
+ from vllm.config.vllm import get_layers_from_vllm_config
225
+ from vllm.model_executor.layers.attention_layer_base import (
226
+ AttentionLayerBase,
227
+ )
228
+ from vllm.utils.math_utils import cdiv
229
+
230
+ cache_config = vllm_config.cache_config
231
+ # special fix for GDN since kernel only supports block size dividable by 64
232
+ attn_layers = get_layers_from_vllm_config(
233
+ vllm_config,
234
+ AttentionLayerBase, # type: ignore[type-abstract]
235
+ )
236
+
237
+ kernel_block_size = None
238
+ for layer in attn_layers.values():
239
+ b = layer.get_attn_backend()
240
+ if b.get_name() == "GDN_ATTN":
241
+ kernel_block_size = 64
242
+ break
243
+
244
+ if kernel_block_size is None:
245
+ return
246
+ new_block_size = (
247
+ cdiv(cache_config.block_size, kernel_block_size) * kernel_block_size
248
+ )
249
+ if new_block_size == cache_config.block_size:
250
+ return
251
+
252
+ if cache_config.mamba_cache_mode == "align":
253
+ cache_config.mamba_block_size = new_block_size
254
+ original_mamba_page_size_padded = cache_config.mamba_page_size_padded
255
+ if cache_config.mamba_page_size_padded is not None:
256
+ attn_page_size_1_token = (
257
+ cache_config.mamba_page_size_padded // cache_config.block_size
258
+ )
259
+ cache_config.mamba_page_size_padded = (
260
+ new_block_size * attn_page_size_1_token
261
+ )
262
+ cache_config.block_size = new_block_size
263
+ logger.info(
264
+ "[XPU]Setting attention block size to %d tokens to ensure multiple of %d, "
265
+ "set mamba_page_size_padded to %d bytes accordingly, before was %d bytes.",
266
+ new_block_size,
267
+ kernel_block_size,
268
+ cache_config.mamba_page_size_padded,
269
+ original_mamba_page_size_padded,
270
+ )
271
+
272
@classmethod
273
def support_hybrid_kv_cache(cls) -> bool:
274
return True