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 4, 17:30 UTC
[XPU] Enable torch.compile for XPU GDN attention (#39466)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com> Signed-off-by: Yuwen Zhou <yuwen.zhou@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Yuwen Zhou committed
Apr 24, 2026 at 16:26 UTC
4c34b2f6fc63435c791c9054c579ca3f8c902bb6
3 files changed
+81
-41
vllm/_xpu_ops.py
+73
@@ -92,6 +92,72 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
92
return torch.empty((M, N), dtype=input.dtype, device=input.device)
93
94
95
+def _gdn_attention_core_xpu_impl(
96
+ core_attn_out: torch.Tensor,
97
+ z: torch.Tensor,
98
+ projected_states_qkvz: torch.Tensor,
99
+ projected_states_ba: torch.Tensor,
100
+ layer_name: str,
101
+) -> None:
102
+ """Custom op wrapping the XPU SYCL GDN kernel for torch.compile."""
103
+ from vllm.forward_context import get_forward_context
104
+ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
105
+
106
+ forward_context = get_forward_context()
107
+ self = forward_context.no_compile_layers[layer_name]
108
+ attn_metadata_raw = forward_context.attn_metadata
109
+
110
+ if attn_metadata_raw is None:
111
+ return
112
+
113
+ assert isinstance(attn_metadata_raw, dict)
114
+ attn_metadata = attn_metadata_raw[self.prefix]
115
+ assert isinstance(attn_metadata, GDNAttentionMetadata)
116
+
117
+ # TODO: xpu does not support speculative decoding yet
118
+ assert attn_metadata.spec_sequence_masks is None # type: ignore[attr-defined]
119
+
120
+ conv_weights = self.conv1d.weight.view(
121
+ self.conv1d.weight.size(0), self.conv1d.weight.size(2)
122
+ )
123
+
124
+ torch.ops._xpu_C.gdn_attention(
125
+ core_attn_out,
126
+ z,
127
+ projected_states_qkvz,
128
+ projected_states_ba,
129
+ self.num_k_heads,
130
+ self.num_v_heads,
131
+ self.head_k_dim,
132
+ self.head_v_dim,
133
+ conv_state=self.kv_cache[0],
134
+ ssm_state=self.kv_cache[1],
135
+ conv_weights=conv_weights,
136
+ conv_bias=self.conv1d.bias,
137
+ activation=self.activation,
138
+ A_log=self.A_log,
139
+ dt_bias=self.dt_bias,
140
+ num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined]
141
+ num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined]
142
+ has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined]
143
+ non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, # type: ignore[attr-defined]
144
+ non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, # type: ignore[attr-defined]
145
+ num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined]
146
+ tp_size=self.tp_size,
147
+ reorder_input=not self.gqa_interleaved_layout,
148
+ )
149
+
150
+
151
+def _gdn_attention_core_xpu_fake(
152
+ core_attn_out: torch.Tensor,
153
+ z: torch.Tensor,
154
+ projected_states_qkvz: torch.Tensor,
155
+ projected_states_ba: torch.Tensor,
156
+ layer_name: str,
157
+) -> None:
158
+ return
159
+
160
+
161
def _xpu_ops_deepseek_scaling_rope_impl(
162
positions: torch.Tensor,
163
query: torch.Tensor,
@@ -618,6 +684,13 @@ class xpu_ops:
684
fake_impl=_xpu_mxfp4_quantize_fake,
685
)
686
687
+ direct_register_custom_op(
688
+ op_name="gdn_attention_core_xpu",
689
+ op_func=_gdn_attention_core_xpu_impl,
690
+ mutates_args=["core_attn_out", "z"],
691
+ fake_impl=_gdn_attention_core_xpu_fake,
692
+ )
693
+
694
_OPS_REGISTERED = True
695
696
vllm/config/compilation.py
+1
@@ -744,6 +744,7 @@ class CompilationConfig:
744
"vllm::linear_attention",
745
"vllm::plamo2_mamba_mixer",
746
"vllm::gdn_attention_core",
747
+ "vllm::gdn_attention_core_xpu",
748
"vllm::olmo_hybrid_gdn_full_forward",
749
"vllm::kda_attention",
750
"vllm::sparse_attn_indexer",
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+7
-41
@@ -618,54 +618,20 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
618
# ============================================================
619
# Part 2: Core Attention
620
# ============================================================
621
- forward_context = get_forward_context()
622
- attn_metadata_raw = forward_context.attn_metadata
621
core_attn_out = torch.zeros(
622
(num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
623
dtype=hidden_states.dtype,
624
device=hidden_states.device,
625
)
626
z = torch.empty_like(core_attn_out)
629
- if attn_metadata_raw is not None:
630
- assert isinstance(attn_metadata_raw, dict)
631
- attn_metadata = attn_metadata_raw[self.prefix]
632
-
633
- # TODO: xpu does not support this param yet
634
- spec_sequence_masks = attn_metadata.spec_sequence_masks # type: ignore[attr-defined]
635
- assert spec_sequence_masks is None
627
637
- conv_weights = self.conv1d.weight.view(
638
- self.conv1d.weight.size(0), self.conv1d.weight.size(2)
639
- )
640
-
641
- conv_state = self.kv_cache[0]
642
- ssm_state = self.kv_cache[1]
643
-
644
- torch.ops._xpu_C.gdn_attention(
645
- core_attn_out,
646
- z,
647
- projected_states_qkvz,
648
- projected_states_ba,
649
- self.num_k_heads,
650
- self.num_v_heads,
651
- self.head_k_dim,
652
- self.head_v_dim,
653
- conv_state=conv_state,
654
- ssm_state=ssm_state,
655
- conv_weights=conv_weights,
656
- conv_bias=self.conv1d.bias,
657
- activation=self.activation,
658
- A_log=self.A_log,
659
- dt_bias=self.dt_bias,
660
- num_prefills=attn_metadata.num_prefills, # type: ignore[attr-defined]
661
- num_decodes=attn_metadata.num_decodes, # type: ignore[attr-defined]
662
- has_initial_state=attn_metadata.has_initial_state, # type: ignore[attr-defined]
663
- non_spec_query_start_loc=attn_metadata.non_spec_query_start_loc, # type: ignore[attr-defined]
664
- non_spec_state_indices_tensor=attn_metadata.non_spec_state_indices_tensor, # type: ignore[attr-defined]
665
- num_actual_tokens=attn_metadata.num_actual_tokens, # type: ignore[attr-defined]
666
- tp_size=self.tp_size,
667
- reorder_input=not self.gqa_interleaved_layout,
668
- )
628
+ torch.ops.vllm.gdn_attention_core_xpu(
629
+ core_attn_out,
630
+ z,
631
+ projected_states_qkvz,
632
+ projected_states_ba,
633
+ self.prefix,
634
+ )
635
636
# ============================================================
637
# Part 3: Output Projection