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 2, 11:00 UTC
[XPU] support MTP of gdn attention (#43565)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Qiming Zhang committed
May 29, 2026 at 02:10 UTC
e8b51999731eab80d88a13da01f08f27400746af
1 file changed
+40
-8
vllm/_xpu_ops.py
+40
-8
@@ -114,8 +114,34 @@ def _gdn_attention_core_xpu_impl(
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]
117
+ num_actual_tokens = attn_metadata.num_actual_tokens
118
+ num_accepted_tokens = attn_metadata.num_accepted_tokens
119
+
120
+ num_prefills = attn_metadata.num_prefills
121
+ num_decodes = attn_metadata.num_decodes
122
+ num_spec_decodes = attn_metadata.num_spec_decodes
123
+
124
+ has_initial_state = attn_metadata.has_initial_state
125
+
126
+ non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
127
+ non_spec_token_indx = attn_metadata.non_spec_token_indx
128
+ non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
129
+ non_spec_state_indices_tensor = (
130
+ non_spec_state_indices_tensor.contiguous()
131
+ if non_spec_state_indices_tensor is not None
132
+ else None
133
+ )
134
+
135
+ spec_query_start_loc = attn_metadata.spec_query_start_loc
136
+ spec_token_indx = attn_metadata.spec_token_indx
137
+ spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
138
+
139
+ spec_sequence_masks = attn_metadata.spec_sequence_masks
140
+ if spec_sequence_masks is not None:
141
+ if non_spec_token_indx is not None:
142
+ non_spec_token_indx = non_spec_token_indx.to(torch.int32)
143
+ if spec_token_indx is not None:
144
+ spec_token_indx = spec_token_indx.to(torch.int32)
145
146
conv_weights = self.conv1d.weight.view(
147
self.conv1d.weight.size(0), self.conv1d.weight.size(2)
@@ -137,12 +163,18 @@ def _gdn_attention_core_xpu_impl(
163
activation=self.activation,
164
A_log=self.A_log,
165
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]
166
+ num_prefills=num_prefills, # type: ignore[attr-defined]
167
+ num_decodes=num_decodes, # type: ignore[attr-defined]
168
+ num_spec_decodes=num_spec_decodes, # type: ignore[attr-defined]
169
+ has_initial_state=has_initial_state, # type: ignore[attr-defined]
170
+ non_spec_query_start_loc=non_spec_query_start_loc, # type: ignore[attr-defined]
171
+ non_spec_token_indx=non_spec_token_indx, # type: ignore[attr-defined]
172
+ non_spec_state_indices_tensor=non_spec_state_indices_tensor, # type: ignore[attr-defined]
173
+ spec_query_start_loc=spec_query_start_loc, # type: ignore[attr-defined]
174
+ spec_token_indx=spec_token_indx, # type: ignore[attr-defined]
175
+ spec_state_indices_tensor=spec_state_indices_tensor,
176
+ num_accepted_tokens=num_accepted_tokens, # type: ignore[attr-defined]
177
+ num_actual_tokens=num_actual_tokens, # type: ignore[attr-defined]
178
tp_size=self.tp_size,
179
reorder_input=not self.gqa_interleaved_layout,
180
)