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 8, 14:13 UTC
[BUGFIX] Fix accuracy bugs in Qwen3-Next MTP (#34077)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Feb 10, 2026 at 19:57 UTC
000214c4bb3f4fb61989eea19c625aedd0559ace
1 file changed
+18
-4
vllm/v1/attention/backends/gdn_attn.py
+18
-4
@@ -208,7 +208,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
208
209
non_spec_query_lens = query_lens[~spec_sequence_masks]
210
num_decodes = (non_spec_query_lens == 1).sum().item()
211
- num_prefills = non_spec_query_lens.size(0) - num_decodes
211
+ # Exclude zero-length padded sequences from prefill count.
212
+ num_zero_len = (non_spec_query_lens == 0).sum().item()
213
+ num_prefills = non_spec_query_lens.size(0) - num_decodes - num_zero_len
214
num_decode_tokens = num_decodes
215
num_prefill_tokens = non_spec_query_lens.sum().item() - num_decode_tokens
216
num_spec_decode_tokens = (
@@ -228,9 +230,15 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
230
non_spec_token_indx = torch.empty(
231
0, dtype=torch.int32, device=query_start_loc.device
232
)
231
- spec_state_indices_tensor = block_table_tensor[:, : self.num_spec + 1]
233
+ # Filter by spec_sequence_masks to exclude padded sequences
234
+ spec_state_indices_tensor = block_table_tensor[
235
+ spec_sequence_masks, : self.num_spec + 1
236
+ ]
237
non_spec_state_indices_tensor = None
233
- spec_query_start_loc = query_start_loc
238
+ # Padded sequences are always at the back, so the first
239
+ # num_spec_decodes + 1 entries of query_start_loc already
240
+ # contain the correct cumulative token counts.
241
+ spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
242
non_spec_query_start_loc = None
243
non_spec_query_start_loc_cpu = None
244
else:
@@ -294,6 +302,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
302
else:
303
has_initial_state = None
304
305
+ # Function code counted on either presency non-spec decode or spec decode,
306
+ # but not both.
307
+ assert not (num_decodes > 0 and num_spec_decodes > 0), (
308
+ f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}"
309
+ )
310
+
311
# Prepare tensors for cudagraph
312
# Note: m.num_actual_tokens is already padded by the model runner for CUDAGraph
313
batch_size = m.num_actual_tokens
@@ -312,7 +326,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
326
spec_state_indices_tensor[num_spec_decodes:].fill_(PAD_SLOT_ID)
327
328
self.spec_sequence_masks[:num_spec_decodes].copy_(
315
- spec_sequence_masks, non_blocking=True
329
+ spec_sequence_masks[:num_spec_decodes], non_blocking=True
330
)
331
spec_sequence_masks = self.spec_sequence_masks[:batch_size]
332
spec_sequence_masks[num_spec_decodes:].fill_(False)