@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f870b92976

[Qwen3.8-Flash-Next] Fuse Qwen4Exp PLE kernels (#54517)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <codex@openai.com>

Thien Tran committed Sep 2, 2026 at 12:32 UTC f870b9297685bd3c968063b52b769315cb08fd7f
6 files changed +1836 -594
tests/models/qwen4_exp/test_ple.py
+983 -65
@@ -1,7 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from dataclasses import dataclass
5 from functools import partial
6 +from itertools import accumulate
7 from types import SimpleNamespace
8
9 import pytest
@@ -17,13 +19,16 @@ from vllm.models.qwen4_exp.common.ple import (
19 compute_ple_shard_overlap,
20 copy_ple_embedding_shard_,
21 )
20 -from vllm.models.qwen4_exp.nvidia import ple_layer as ple_layer_module
22 from vllm.models.qwen4_exp.nvidia.ple_layer import (
23 Qwen4ExpNGramEmbedding,
24 Qwen4ExpPLEFp8EmbeddingMethod,
25 Qwen4ExpPLELayer,
26 _get_ple_embedding_quant_method,
27 )
28 +from vllm.v1.attention.backends.short_conv_attn import (
29 + PleShortConvAttentionMetadata,
30 +)
31 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
32
33
34 def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
@@ -283,54 +288,9 @@ def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
288 assert _get_ple_embedding_quant_method(quant_config, prefix) is None
289
290
286 -def test_ple_ngram_ids_custom_op_uses_current_request_layout(monkeypatch) -> None:
287 - class RuntimeNGramEmbedding(nn.Module):
288 - def compute_ngram_ids(
289 - self,
290 - input_ids: torch.Tensor,
291 - query_start_loc: torch.Tensor,
292 - ngram_context: torch.Tensor,
293 - ) -> torch.Tensor:
294 - del input_ids, ngram_context
295 - num_reqs = query_start_loc.numel() - 1
296 - return torch.full((4, 2), num_reqs, dtype=torch.long)
297 -
298 - layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
299 - nn.Module.__init__(layer)
300 - layer.ple_embedding = RuntimeNGramEmbedding()
301 - monkeypatch.setattr(
302 - ple_layer_module,
303 - "get_forward_context",
304 - lambda: SimpleNamespace(no_compile_layers={"ple": layer}),
305 - )
306 - input_ids = torch.arange(4)
307 - ngram_context = torch.zeros(2, 2, dtype=torch.long)
308 - output = torch.empty(4, 2, dtype=torch.long)
309 -
310 - ple_layer_module.qwen4_exp_compute_ple_ngram_ids(
311 - input_ids,
312 - torch.tensor([0, 4]),
313 - ngram_context,
314 - output,
315 - "ple",
316 - )
317 - assert torch.equal(output, torch.ones_like(output))
318 -
319 - ple_layer_module.qwen4_exp_compute_ple_ngram_ids(
320 - input_ids,
321 - torch.tensor([0, 2, 4]),
322 - ngram_context,
323 - output,
324 - "ple",
325 - )
326 - assert torch.equal(output, torch.full_like(output, 2))
327 -
328 -
291 def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
330 - module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
331 - nn.Module.__init__(module)
332 - module.conv_state_len = 6
333 - module.short_conv_dilation = 2
292 + conv_state_len = 6
293 + dilation = 2
294
295 conv_weights = torch.tensor([[0.25, -0.5, 0.75, 1.0]])
296 conv_state = torch.zeros(2, 1, 9)
@@ -338,12 +298,12 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
298 first_inputs = torch.tensor([[10.0], [20.0], [30.0], [40.0]])
299 initial_state = conv_state[1:].clone()
300 first_history = torch.cat(
341 - (initial_state[..., : module.conv_state_len], first_inputs.T.unsqueeze(0)),
301 + (initial_state[..., :conv_state_len], first_inputs.T.unsqueeze(0)),
302 dim=-1,
303 )
304
305 graph_padded_inputs = F.pad(first_inputs, (0, 0, 0, 4))
346 - first_output = module._short_conv_dilated_spec_batched(
306 + first_output = _short_conv_dilated_spec_pytorch(
307 graph_padded_inputs,
308 conv_state,
309 conv_weights,
@@ -351,6 +311,8 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
311 torch.tensor([0, 4, 4]),
312 torch.tensor([1, 0]),
313 spec_query_len=4,
314 + conv_state_len=conv_state_len,
315 + dilation=dilation,
316 )
317
318 expected_first_output = F.silu(
@@ -358,7 +320,7 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
320 first_history,
321 conv_weights.unsqueeze(1),
322 groups=1,
361 - dilation=module.short_conv_dilation,
323 + dilation=dilation,
324 )
325 ).transpose(1, 2)[0]
326 expected_first_state = first_history[..., 1:10]
@@ -374,7 +336,7 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
336 expected_second_state = expected_first_state.clone()
337 expected_second_state[..., :7] = second_history[..., 1:8]
338
377 - second_output = module._short_conv_dilated_spec_batched(
339 + second_output = _short_conv_dilated_spec_pytorch(
340 second_inputs,
341 conv_state,
342 conv_weights,
@@ -382,6 +344,8 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
344 torch.tensor([0, 2]),
345 torch.tensor([2]),
346 spec_query_len=4,
347 + conv_state_len=conv_state_len,
348 + dilation=dilation,
349 )
350
351 expected_second_output = F.silu(
@@ -389,28 +353,982 @@ def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
353 second_history,
354 conv_weights.unsqueeze(1),
355 groups=1,
392 - dilation=module.short_conv_dilation,
356 + dilation=dilation,
357 )
358 ).transpose(1, 2)[0, :2]
359 torch.testing.assert_close(second_output, expected_second_output)
360 torch.testing.assert_close(conv_state[1:], expected_second_state)
361
362
399 -def test_ple_short_conv_uses_fallback_when_profile_metadata_is_omitted(
400 - monkeypatch: pytest.MonkeyPatch,
363 +def _reference_ngram_ids(
364 + input_ids: torch.Tensor,
365 + query_start_loc: torch.Tensor,
366 + ngram_context: torch.Tensor,
367 + layer_multipliers: torch.Tensor,
368 + ngram_heads_vocab_sizes: torch.Tensor,
369 + ngram_heads_offsets: torch.Tensor,
370 + eos_token_id: int,
371 + heads_per_ngram: int,
372 +) -> torch.Tensor:
373 + tokens = input_ids.tolist()
374 + starts = query_start_loc.tolist()
375 + contexts = ngram_context.tolist()
376 + multipliers = layer_multipliers.cpu()
377 + sizes = ngram_heads_vocab_sizes.cpu()
378 + offsets = ngram_heads_offsets.cpu()
379 + rows = []
380 + context_len = len(contexts[0])
381 + for req, (start, end) in enumerate(zip(starts, starts[1:])):
382 + history = contexts[req] + tokens[start:end]
383 + for pos in range(context_len, len(history)):
384 + shifted = [history[pos]]
385 + crossed_eos = False
386 + for shift in range(1, context_len + 1):
387 + token = eos_token_id if crossed_eos else history[pos - shift]
388 + shifted.append(token)
389 + crossed_eos |= token == eos_token_id
390 +
391 + row = []
392 + mixed = torch.tensor(shifted[0], dtype=torch.int64) * multipliers[0]
393 + for ngram_order in range(2, context_len + 2):
394 + shift = ngram_order - 1
395 + mixed ^= shifted[shift] * multipliers[shift]
396 + head_start = (ngram_order - 2) * heads_per_ngram
397 + for head in range(head_start, head_start + heads_per_ngram):
398 + row.append(torch.remainder(mixed, sizes[head]) + offsets[head])
399 + rows.append(torch.stack(row))
400 + return torch.stack(rows).to(input_ids.device)
401 +
402 +
403 +_NGRAM_MULTIPLIERS = (
404 + 18_014_398_509_481_983,
405 + 17_114_398_509_481_981,
406 + 16_214_398_509_481_979,
407 + 15_314_398_509_481_977,
408 +)
409 +_NGRAM_HEADS_VOCAB_SIZES = (
410 + 101,
411 + 103,
412 + 107,
413 + 109,
414 + 113,
415 + 127,
416 + 131,
417 + 137,
418 + 139,
419 + 149,
420 + 151,
421 + 157,
422 + 163,
423 + 167,
424 + 173,
425 + 179,
426 + 181,
427 + 191,
428 + 193,
429 + 197,
430 + 199,
431 + 211,
432 + 223,
433 + 227,
434 +)
435 +_NGRAM_EOS_TOKEN_ID = 251
436 +_NGRAM_HEADS_PER_NGRAM = 8
437 +
438 +
439 +def _ngram_hash_params(device: torch.device, context_len: int) -> dict:
440 + num_heads = context_len * _NGRAM_HEADS_PER_NGRAM
441 + sizes = torch.tensor(
442 + _NGRAM_HEADS_VOCAB_SIZES[:num_heads], dtype=torch.int64, device=device
443 + )
444 + offsets = torch.zeros_like(sizes)
445 + offsets[1:] = torch.cumsum(sizes, dim=0)[:-1]
446 + return {
447 + "layer_multipliers": torch.tensor(
448 + _NGRAM_MULTIPLIERS[: context_len + 1],
449 + dtype=torch.int64,
450 + device=device,
451 + ),
452 + "ngram_heads_vocab_sizes": sizes,
453 + "ngram_heads_offsets": offsets,
454 + "eos_token_id": _NGRAM_EOS_TOKEN_ID,
455 + "heads_per_ngram": _NGRAM_HEADS_PER_NGRAM,
456 + }
457 +
458 +
459 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused PLE needs CUDA")
460 +@pytest.mark.parametrize(
461 + ("query_lens", "eos_offsets", "contexts", "first_token_id"),
462 + [
463 + ([1], [], [[11, 12]], 20),
464 + ([3], [1], [[11]], 20),
465 + ([4, 4], [0, 7], [[11, 12], [13, 14]], 20),
466 + ([4, 0, 3], [], [[11, 12], [13, 14], [15, 16]], 20),
467 + (
468 + [1, 33, 2],
469 + [5, 32],
470 + [[_NGRAM_EOS_TOKEN_ID, 11], [12, 13], [14, _NGRAM_EOS_TOKEN_ID]],
471 + 20,
472 + ),
473 + (
474 + [5, 12, 16, 1, 16, 17],
475 + [10, 40],
476 + [[11, 12], [13, 14], [15, 16], [17, 18], [19, 20], [21, 22]],
477 + 20,
478 + ),
479 + ([5, 3], [3], [[11, 12, 13], [14, 15, 16]], 20),
480 + (
481 + [4, 0, 3],
482 + [],
483 + [
484 + [200_000, 200_001],
485 + [250_000, 250_001],
486 + [300_000, 300_001],
487 + ],
488 + 350_000,
489 + ),
490 + ([3], [], [[1_000_000_000, 1_000_000_001]], 1_000_000_002),
491 + ],
492 + ids=[
493 + "single-token",
494 + "bigram-only",
495 + "power-of-two",
496 + "empty-request",
497 + "three-requests",
498 + "six-requests",
499 + "four-gram",
500 + "int64-overflow",
501 + "large-int32-ids",
502 + ],
503 +)
504 +def test_fused_ngram_ids_correctness(
505 + query_lens: list[int],
506 + eos_offsets: list[int],
507 + contexts: list[list[int]],
508 + first_token_id: int,
509 ) -> None:
510 + from vllm.models.qwen4_exp.nvidia.ops.ple import ple_ngram_ids
511 +
512 + device = torch.device("cuda")
513 + query_start_loc = torch.tensor(
514 + [0, *accumulate(query_lens)], dtype=torch.int32, device=device
515 + )
516 + input_ids = torch.arange(
517 + first_token_id,
518 + first_token_id + sum(query_lens),
519 + dtype=torch.int32,
520 + device=device,
521 + )
522 + if eos_offsets:
523 + input_ids[eos_offsets] = _NGRAM_EOS_TOKEN_ID
524 + ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device)
525 + params = _ngram_hash_params(device, ngram_context.shape[1])
526 +
527 + expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
528 + actual = ple_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
529 +
530 + assert torch.equal(actual, expected)
531 + offsets = params["ngram_heads_offsets"]
532 + sizes = params["ngram_heads_vocab_sizes"]
533 + assert torch.all((actual >= offsets) & (actual < offsets + sizes))
534 +
535 +
536 +def _short_conv_dilated_decode_pytorch(
537 + x_d: torch.Tensor,
538 + conv_state: torch.Tensor,
539 + conv_weights: torch.Tensor,
540 + state_indices_tensor_d: torch.Tensor,
541 + has_initial_states_d: torch.Tensor | None,
542 + *,
543 + conv_state_len: int,
544 + dilation: int,
545 +) -> torch.Tensor:
546 + state_indices = state_indices_tensor_d.to(
547 + device=conv_state.device, dtype=torch.int64
548 + )
549 + # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
550 + # slot 0 for a safe gather, then zero output and skip write-back.
551 + valid_state = state_indices != NULL_BLOCK_ID
552 + state_indices = torch.where(
553 + valid_state, state_indices, torch.zeros_like(state_indices)
554 + )
555 + if has_initial_states_d is None:
556 + has_initial_state = valid_state
557 + else:
558 + if has_initial_states_d.numel() < state_indices_tensor_d.numel():
559 + raise ValueError(
560 + "has_initial_states_d size mismatch: "
561 + f"got {has_initial_states_d.numel()}, "
562 + f"need >= {state_indices_tensor_d.numel()}."
563 + )
564 + has_initial_state = has_initial_states_d[: state_indices_tensor_d.numel()].to(
565 + device=conv_state.device, dtype=torch.bool
566 + )
567 + has_initial_state = has_initial_state & valid_state
568 +
569 + cached_state = conv_state.index_select(0, state_indices)
570 + state = cached_state[..., :conv_state_len].to(x_d.dtype)
571 + if conv_state_len > 0:
572 + initial_state = torch.where(
573 + has_initial_state.view(-1, 1, 1),
574 + state,
575 + torch.zeros_like(state),
576 + )
577 + history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
578 + else:
579 + history = x_d.unsqueeze(-1)
580 +
581 + conv_output = F.conv1d(
582 + history,
583 + conv_weights.unsqueeze(1).contiguous(),
584 + groups=history.size(1),
585 + dilation=dilation,
586 + ).squeeze(-1)
587 + output = F.silu(conv_output)
588 + output = output * valid_state.view(-1, 1).to(output.dtype)
589 +
590 + if conv_state_len > 0:
591 + next_state = history[..., -conv_state_len:]
592 + # Padded rows are remapped to the reserved null slot. Preserve its
593 + # existing value while writing the new states for valid rows.
594 + existing_base_state = cached_state[..., :conv_state_len]
595 + safe_next_state = torch.where(
596 + valid_state.view(-1, 1, 1),
597 + next_state.to(conv_state.dtype),
598 + existing_base_state,
599 + )
600 + cached_state[..., :conv_state_len] = safe_next_state
601 + conv_state.index_copy_(0, state_indices, cached_state)
602 +
603 + return output
604 +
605 +
606 +def _short_conv_dilated_prefill_pytorch(
607 + x_p: torch.Tensor,
608 + metadata: PleShortConvAttentionMetadata,
609 + conv_state: torch.Tensor,
610 + conv_weights: torch.Tensor,
611 + state_indices_tensor_p: torch.Tensor,
612 + num_prefills: int,
613 + num_decode_tokens: int,
614 + num_prefill_tokens: int,
615 + *,
616 + conv_state_len: int,
617 + dilation: int,
618 +) -> torch.Tensor:
619 + # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
620 + # requests and equals ``query_start_loc`` when spec-decode is inactive.
621 + non_spec_query_start_loc = metadata.non_spec_query_start_loc
622 + if non_spec_query_start_loc is None:
623 + raise ValueError("query_start_loc is required for prefill short-conv")
624 + query_start_loc_p = (
625 + non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
626 + )
627 + # The metadata builder guarantees that the prefill query offsets start
628 + # at 0 and end at num_prefill_tokens. Avoid reading those values here,
629 + # since doing so would force a device-to-host synchronization.
630 + has_initial_states_p = metadata.has_initial_states_p
631 + if has_initial_states_p is None:
632 + raise ValueError("has_initial_states_p is required for prefill short-conv")
633 +
634 + output = torch.empty_like(x_p)
635 + q_starts = query_start_loc_p.to(torch.int64)
636 + if state_indices_tensor_p.numel() < num_prefills:
637 + raise ValueError(
638 + "state_indices_tensor_p size mismatch: "
639 + f"got {state_indices_tensor_p.numel()}, "
640 + f"need >= {num_prefills}."
641 + )
642 + if has_initial_states_p.numel() < num_prefills:
643 + raise ValueError(
644 + "has_initial_states_p size mismatch: "
645 + f"got {has_initial_states_p.numel()}, "
646 + f"need >= {num_prefills}."
647 + )
648 + if num_prefills == 0 or x_p.numel() == 0:
649 + return output
650 + lengths = q_starts[1:] - q_starts[:-1]
651 + # Use the CPU-computed packing width from the metadata builder instead
652 + # of synchronizing on lengths.max().
653 + max_len = metadata.max_prefill_query_len
654 + if max_len <= 0:
655 + return output
656 +
657 + hidden_size = x_p.shape[1]
658 + positions = torch.arange(num_prefill_tokens, device=x_p.device, dtype=torch.int64)
659 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
660 + col_indices = positions - q_starts[req_indices]
661 +
662 + packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
663 + packed_tokens[req_indices, col_indices] = x_p
664 + packed_tokens = packed_tokens.transpose(1, 2).contiguous()
665 +
666 + state_indices = state_indices_tensor_p[:num_prefills].to(
667 + device=conv_state.device, dtype=torch.int64
668 + )
669 + valid_state = state_indices != NULL_BLOCK_ID
670 + state_indices = torch.where(
671 + valid_state, state_indices, torch.zeros_like(state_indices)
672 + )
673 + has_initial = has_initial_states_p[:num_prefills].to(
674 + device=conv_state.device, dtype=torch.bool
675 + )
676 + if conv_state_len > 0:
677 + if conv_state.shape[0] == 0:
678 + state = conv_state.new_zeros(
679 + (num_prefills, hidden_size, conv_state_len),
680 + dtype=x_p.dtype,
681 + )
682 + else:
683 + state = conv_state.index_select(0, state_indices)[..., :conv_state_len].to(
684 + x_p.dtype
685 + )
686 + use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
687 + initial_state = torch.where(
688 + use_initial_mask,
689 + state,
690 + torch.zeros_like(state),
691 + )
692 + history = torch.cat((initial_state, packed_tokens), dim=-1)
693 + else:
694 + history = packed_tokens
695 +
696 + conv_output = F.conv1d(
697 + history,
698 + conv_weights.unsqueeze(1).contiguous(),
699 + groups=history.size(1),
700 + dilation=dilation,
701 + )
702 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
703 +
704 + token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
705 + valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
706 + valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
707 + num_prefills, 1
708 + )
709 + conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
710 + output.copy_(conv_output[req_indices, col_indices])
711 +
712 + if conv_state_len > 0 and conv_state.shape[0] > 0:
713 + state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
714 + num_prefills, 1, 1
715 + )
716 + state_offsets = torch.arange(
717 + conv_state_len, device=history.device, dtype=torch.int64
718 + ).view(1, 1, conv_state_len)
719 + next_state = history.gather(
720 + dim=2,
721 + index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
722 + )
723 + # Write back without a host synchronization. Valid, non-empty rows
724 + # receive their new state; padding and zero-length rows keep the
725 + # current cache value.
726 + existing_state = conv_state.index_select(0, state_indices)
727 + existing_base_state = existing_state[..., :conv_state_len]
728 + update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
729 + safe_next_state = torch.where(
730 + update_mask.view(num_prefills, 1, 1),
731 + next_state.to(conv_state.dtype),
732 + existing_base_state,
733 + )
734 + existing_state[..., :conv_state_len] = safe_next_state
735 + conv_state.index_copy_(0, state_indices, existing_state)
736 + return output
737 +
738 +
739 +def _short_conv_dilated_spec_pytorch(
740 + x_spec: torch.Tensor,
741 + conv_state: torch.Tensor,
742 + conv_weights: torch.Tensor,
743 + spec_state_indices_tensor: torch.Tensor,
744 + spec_query_start_loc: torch.Tensor,
745 + num_accepted_tokens: torch.Tensor,
746 + spec_query_len: int,
747 + *,
748 + conv_state_len: int,
749 + dilation: int,
750 +) -> torch.Tensor:
751 + """Dilated short-conv for speculative-decode (MTP) requests.
752 +
753 + Each spec request feeds multiple (draft + 1) query tokens. The conv
754 + outputs are computed causally after rolling back the previous draft
755 + state by ``num_accepted_tokens - 1``. The current candidate inputs stay
756 + in the extended cache for the next forward, matching
757 + ``causal_conv1d_update``.
758 +
759 + ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
760 + length and is a Python int, so no host synchronization is needed; this
761 + keeps the path safe for full CUDA-graph capture/replay where the buffers
762 + are padded at the request level.
763 + """
764 + num_reqs = spec_state_indices_tensor.numel()
765 + hidden_size = x_spec.size(-1)
766 + # Use a fixed packing width instead of synchronizing on lengths.max().
767 + max_len = spec_query_len
768 + # Full CUDA graphs can pad these buffers. Only the first num_reqs
769 + # accepted-token counts belong to actual speculative requests.
770 + num_accepted_tokens = num_accepted_tokens[:num_reqs]
771 + q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
772 + # Keep the number of real speculative tokens on the device.
773 + total_real_tokens = q_starts[num_reqs]
774 +
775 + state_indices = spec_state_indices_tensor.to(
776 + device=conv_state.device, dtype=torch.int64
777 + )
778 + valid_state = state_indices != NULL_BLOCK_ID
779 + state_indices = torch.where(
780 + valid_state, state_indices, torch.zeros_like(state_indices)
781 + )
782 + positions = torch.arange(x_spec.size(0), device=x_spec.device, dtype=torch.int64)
783 + # Route graph-padded token rows to the discarded dummy request so that
784 + # they cannot overwrite real packed data.
785 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
786 + valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
787 + clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
788 + col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
789 + pack_req_indices = torch.where(
790 + valid_tokens,
791 + clamped_req_indices,
792 + torch.full_like(req_indices, num_reqs),
793 + )
794 + pack_col_indices = torch.where(
795 + valid_tokens, col_indices, torch.zeros_like(col_indices)
796 + )
797 +
798 + # The last request row is the dummy sink for graph padding.
799 + packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
800 + packed[pack_req_indices, pack_col_indices] = x_spec
801 + packed = packed.transpose(1, 2).contiguous()
802 +
803 + if conv_state_len > 0:
804 + cached_state = conv_state.index_select(0, state_indices)
805 + rollback_offsets = num_accepted_tokens.to(
806 + device=conv_state.device, dtype=torch.int64
807 + ).sub(1)
808 + rollback_offsets = torch.where(
809 + valid_state,
810 + rollback_offsets.clamp_(0, max_len - 1),
811 + torch.zeros_like(rollback_offsets),
812 + )
813 + state_offsets = torch.arange(
814 + conv_state_len, device=conv_state.device, dtype=torch.int64
815 + ).view(1, 1, conv_state_len)
816 + rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
817 + state = cached_state.gather(2, rollback_indices.expand(-1, hidden_size, -1)).to(
818 + x_spec.dtype
819 + )
820 + state = torch.where(
821 + valid_state.view(num_reqs, 1, 1),
822 + state,
823 + torch.zeros_like(state),
824 + )
825 + # Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
826 + dummy_state = state.new_zeros((1, hidden_size, conv_state_len))
827 + state_full = torch.cat((state, dummy_state), dim=0)
828 + history = torch.cat((state_full, packed), dim=-1)
829 + else:
830 + history = packed
831 +
832 + conv_output = F.conv1d(
833 + history,
834 + conv_weights.unsqueeze(1).contiguous(),
835 + groups=history.size(1),
836 + dilation=dilation,
837 + )
838 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
839 +
840 + output = conv_output[pack_req_indices, pack_col_indices]
841 + output = output * valid_tokens.view(-1, 1).to(output.dtype)
842 +
843 + # Keep all current candidate inputs in the extended state. On the next
844 + # target forward, ``num_accepted_tokens - 1`` selects the rollback
845 + # window before processing the newly scheduled tokens.
846 + if conv_state_len > 0:
847 + state_capacity = conv_state_len + max_len - 1
848 + if conv_state.size(-1) < state_capacity:
849 + raise RuntimeError(
850 + "PLE short-conv cache cannot retain speculative tokens: "
851 + f"got {conv_state.size(-1)}, need {state_capacity}."
852 + )
853 + candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
854 + query_lengths = q_starts[1:] - q_starts[:-1]
855 + state_positions = torch.arange(
856 + state_capacity, device=history.device, dtype=torch.int64
857 + ).view(1, 1, state_capacity)
858 + update_lengths = (conv_state_len + query_lengths - 1).view(num_reqs, 1, 1)
859 + update_mask = valid_state.view(num_reqs, 1, 1) & (
860 + state_positions < update_lengths
861 + )
862 + existing_state = cached_state[..., :state_capacity]
863 + next_state = torch.where(
864 + update_mask,
865 + candidate_state.to(conv_state.dtype),
866 + existing_state,
867 + )
868 + cached_state[..., :state_capacity] = next_state
869 + conv_state.index_copy_(0, state_indices, cached_state)
870 +
871 + return output
872 +
873 +
874 +def _short_conv_dilated_dispatch_pytorch(
875 + inputs: torch.Tensor,
876 + residual: torch.Tensor,
877 + metadata: PleShortConvAttentionMetadata,
878 + conv_state: torch.Tensor,
879 + conv_weights: torch.Tensor,
880 + *,
881 + conv_state_len: int,
882 + dilation: int,
883 +) -> None:
884 + """PyTorch reference for fused short-convolution dispatch."""
885 + num_prefills = metadata.num_prefills
886 + num_decodes = metadata.num_decodes
887 + num_decode_tokens = metadata.num_decode_tokens
888 + num_prefill_tokens = metadata.num_prefill_tokens
889 + has_prefill = num_prefills > 0
890 + has_decode = num_decodes > 0
891 + has_spec = metadata.spec_sequence_masks is not None
892 + x = inputs[: metadata.num_actual_tokens]
893 + residual = residual[: metadata.num_actual_tokens]
894 +
895 + if has_spec:
896 + if has_prefill or has_decode:
897 + assert metadata.spec_token_indx is not None
898 + assert metadata.non_spec_token_indx is not None
899 + spec_token_indices = metadata.spec_token_indx.long()
900 + non_spec_token_indices = metadata.non_spec_token_indx.long()
901 + x_spec = x.index_select(0, spec_token_indices)
902 + x_non_spec = x.index_select(0, non_spec_token_indices)
903 + residual_spec = residual.index_select(0, spec_token_indices)
904 + residual_non_spec = residual.index_select(0, non_spec_token_indices)
905 + else:
906 + spec_token_indices = None
907 + non_spec_token_indices = None
908 + x_spec = x
909 + x_non_spec = None
910 + residual_spec = residual
911 + residual_non_spec = None
912 + else:
913 + spec_token_indices = None
914 + non_spec_token_indices = None
915 + x_spec = None
916 + x_non_spec = x
917 + residual_spec = None
918 + residual_non_spec = residual
919 +
920 + if has_spec:
921 + assert metadata.spec_state_indices_tensor is not None
922 + assert metadata.spec_query_start_loc is not None
923 + assert metadata.num_accepted_tokens is not None
924 + assert x_spec is not None
925 + assert residual_spec is not None
926 + spec_state_indices = metadata.spec_state_indices_tensor[
927 + : metadata.num_spec_decodes
928 + ]
929 + conv_output = _short_conv_dilated_spec_pytorch(
930 + x_spec=x_spec,
931 + conv_state=conv_state,
932 + conv_weights=conv_weights,
933 + spec_state_indices_tensor=spec_state_indices,
934 + spec_query_start_loc=metadata.spec_query_start_loc,
935 + num_accepted_tokens=metadata.num_accepted_tokens,
936 + spec_query_len=metadata.spec_query_len,
937 + conv_state_len=conv_state_len,
938 + dilation=dilation,
939 + )
940 + residual_spec.add_(conv_output)
941 +
942 + state_indices = metadata.state_indices_tensor
943 + if x_non_spec is not None:
944 + assert state_indices is not None
945 + assert residual_non_spec is not None
946 + if has_prefill:
947 + state_indices_d, state_indices_p = torch.split(
948 + state_indices, [num_decodes, num_prefills], dim=0
949 + )
950 + x_d, x_p = torch.split(
951 + x_non_spec, [num_decode_tokens, num_prefill_tokens], dim=0
952 + )
953 + residual_d, residual_p = torch.split(
954 + residual_non_spec,
955 + [num_decode_tokens, num_prefill_tokens],
956 + dim=0,
957 + )
958 + if has_decode:
959 + conv_output = _short_conv_dilated_decode_pytorch(
960 + x_d=x_d,
961 + conv_state=conv_state,
962 + conv_weights=conv_weights,
963 + state_indices_tensor_d=state_indices_d,
964 + has_initial_states_d=metadata.has_initial_states_d,
965 + conv_state_len=conv_state_len,
966 + dilation=dilation,
967 + )
968 + residual_d.add_(conv_output)
969 + conv_output = _short_conv_dilated_prefill_pytorch(
970 + x_p=x_p,
971 + metadata=metadata,
972 + conv_state=conv_state,
973 + conv_weights=conv_weights,
974 + state_indices_tensor_p=state_indices_p,
975 + num_prefills=num_prefills,
976 + num_decode_tokens=num_decode_tokens,
977 + num_prefill_tokens=num_prefill_tokens,
978 + conv_state_len=conv_state_len,
979 + dilation=dilation,
980 + )
981 + residual_p.add_(conv_output)
982 + else:
983 + conv_output = _short_conv_dilated_decode_pytorch(
984 + x_d=x_non_spec,
985 + conv_state=conv_state,
986 + conv_weights=conv_weights,
987 + state_indices_tensor_d=state_indices[: x_non_spec.size(0)],
988 + has_initial_states_d=metadata.has_initial_states_d,
989 + conv_state_len=conv_state_len,
990 + dilation=dilation,
991 + )
992 + residual_non_spec.add_(conv_output)
993 +
994 + if has_spec and residual_non_spec is not None:
995 + assert spec_token_indices is not None
996 + assert non_spec_token_indices is not None
997 + assert residual_spec is not None
998 + residual.index_copy_(0, spec_token_indices, residual_spec)
999 + residual.index_copy_(0, non_spec_token_indices, residual_non_spec)
1000 +
1001 +
1002 +_KERNEL_STATE_BLOCKS = 64
1003 +
1004 +
1005 +def _make_conv_case(
1006 + device: torch.device,
1007 + seed: int,
1008 + channels: int,
1009 + kernel_size: int,
1010 + dilation: int,
1011 + state_layout: str,
1012 + spec_query_len: int = 1,
1013 +) -> tuple[torch.Generator, torch.Tensor, torch.Tensor, torch.Tensor]:
1014 + rng = torch.Generator(device=device).manual_seed(seed)
1015 + state_len = (kernel_size - 1) * dilation
1016 + state_width = state_len + spec_query_len - 1
1017 + state_shape = (
1018 + (_KERNEL_STATE_BLOCKS, state_width, channels)
1019 + if state_layout == "SD"
1020 + else (_KERNEL_STATE_BLOCKS, channels, state_width)
1021 + )
1022 + state_kernel = torch.randn(
1023 + state_shape, device=device, dtype=torch.bfloat16, generator=rng
1024 + )
1025 + conv_state = (
1026 + state_kernel.transpose(-1, -2) if state_layout == "SD" else state_kernel
1027 + )
1028 + state_reference = conv_state.clone()
1029 + weights = torch.randn(
1030 + channels, kernel_size, device=device, dtype=torch.bfloat16, generator=rng
1031 + )
1032 + return rng, state_reference, conv_state, weights
1033 +
1034 +
1035 +@dataclass(frozen=True)
1036 +class _ConvBatchCase:
1037 + spec_query_lens: tuple[int, ...] = ()
1038 + num_accepted: tuple[int, ...] = ()
1039 + num_decodes: int = 0
1040 + prefill_query_lens: tuple[int, ...] = ()
1041 + channels: int = 2048
1042 + kernel_size: int = 4
1043 + dilation: int = 3
1044 + spec_query_len: int = 1
1045 + graph_padding: int = 0
1046 +
1047 +
1048 +def _make_conv_metadata(
1049 + case: _ConvBatchCase,
1050 + device: torch.device,
1051 +) -> tuple[SimpleNamespace, int]:
1052 + request_groups: dict[str, list[int]] = {
1053 + "spec": [],
1054 + "decode": [],
1055 + "prefill": [],
1056 + }
1057 + token_offset = 0
1058 + max_num_reqs = max(
1059 + len(case.spec_query_lens),
1060 + case.num_decodes,
1061 + len(case.prefill_query_lens),
1062 + )
1063 + request_kinds = []
1064 + for req_idx in range(max_num_reqs):
1065 + if req_idx < len(case.spec_query_lens):
1066 + request_kinds.append(("spec", case.spec_query_lens[req_idx]))
1067 + if req_idx < case.num_decodes:
1068 + request_kinds.append(("decode", 1))
1069 + if req_idx < len(case.prefill_query_lens):
1070 + request_kinds.append(("prefill", case.prefill_query_lens[req_idx]))
1071 + for kind, query_len in request_kinds:
1072 + request_groups[kind].extend(range(token_offset, token_offset + query_len))
1073 + token_offset += query_len
1074 +
1075 + num_spec_reqs = len(case.spec_query_lens)
1076 + num_prefills = len(case.prefill_query_lens)
1077 + has_spec = num_spec_reqs > 0
1078 + has_non_spec = case.num_decodes > 0 or num_prefills > 0
1079 + mixed_batch = has_spec and has_non_spec
1080 +
1081 + spec_state_indices = torch.arange(
1082 + 1, num_spec_reqs + 1, dtype=torch.int32, device=device
1083 + )
1084 + if num_spec_reqs >= 3:
1085 + spec_state_indices[1] = NULL_BLOCK_ID
1086 + non_spec_state_indices = torch.arange(
1087 + num_spec_reqs + 1,
1088 + num_spec_reqs + case.num_decodes + num_prefills + 1,
1089 + dtype=torch.int32,
1090 + device=device,
1091 + )
1092 + if case.num_decodes > 1:
1093 + non_spec_state_indices[case.num_decodes - 1] = NULL_BLOCK_ID
1094 + if num_prefills > 1:
1095 + non_spec_state_indices[case.num_decodes + 1] = NULL_BLOCK_ID
1096 +
1097 + spec_query_start_loc = torch.tensor(
1098 + [0, *accumulate(case.spec_query_lens)], dtype=torch.int32, device=device
1099 + )
1100 + non_spec_query_lens = (1,) * case.num_decodes + case.prefill_query_lens
1101 + non_spec_query_start_loc = torch.tensor(
1102 + [0, *accumulate(non_spec_query_lens)], dtype=torch.int32, device=device
1103 + )
1104 + spec_token_indices = torch.tensor(
1105 + request_groups["spec"], dtype=torch.int64, device=device
1106 + )
1107 + non_spec_token_indices = torch.tensor(
1108 + request_groups["decode"] + request_groups["prefill"],
1109 + dtype=torch.int64,
1110 + device=device,
1111 + )
1112 +
1113 + metadata = SimpleNamespace(
1114 + num_prefills=num_prefills,
1115 + num_decodes=case.num_decodes,
1116 + num_decode_tokens=case.num_decodes,
1117 + num_prefill_tokens=sum(case.prefill_query_lens),
1118 + spec_sequence_masks=(
1119 + torch.tensor([kind == "spec" for kind, _ in request_kinds], device=device)
1120 + if has_spec
1121 + else None
1122 + ),
1123 + spec_token_indx=spec_token_indices if mixed_batch else None,
1124 + non_spec_token_indx=non_spec_token_indices if mixed_batch else None,
1125 + num_actual_tokens=token_offset + case.graph_padding,
1126 + spec_state_indices_tensor=spec_state_indices if has_spec else None,
1127 + spec_query_start_loc=spec_query_start_loc if has_spec else None,
1128 + num_accepted_tokens=(
1129 + torch.tensor(case.num_accepted, dtype=torch.int32, device=device)
1130 + if has_spec
1131 + else None
1132 + ),
1133 + spec_query_len=case.spec_query_len,
1134 + state_indices_tensor=non_spec_state_indices if has_non_spec else None,
1135 + has_initial_states_d=(
1136 + torch.arange(case.num_decodes, device=device) % 2 == 0
1137 + if case.num_decodes
1138 + else None
1139 + ),
1140 + non_spec_query_start_loc=(non_spec_query_start_loc if has_non_spec else None),
1141 + has_initial_states_p=(
1142 + torch.arange(num_prefills, device=device) % 2 == 0 if num_prefills else None
1143 + ),
1144 + max_prefill_query_len=max(case.prefill_query_lens, default=0),
1145 + num_spec_decodes=num_spec_reqs,
1146 + )
1147 + return metadata, token_offset
1148 +
1149 +
1150 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused conv needs CUDA")
1151 +@pytest.mark.parametrize("state_layout", ["SD", "DS"])
1152 +@pytest.mark.parametrize(
1153 + "case",
1154 + [
1155 + pytest.param(
1156 + _ConvBatchCase(num_decodes=1, channels=512),
1157 + id="decode-single",
1158 + ),
1159 + pytest.param(
1160 + _ConvBatchCase(
1161 + num_decodes=33,
1162 + channels=2048,
1163 + kernel_size=5,
1164 + dilation=1,
1165 + ),
1166 + id="decode-padded",
1167 + ),
1168 + pytest.param(
1169 + _ConvBatchCase(prefill_query_lens=(37, 0, 5, 128)),
1170 + id="prefill",
1171 + ),
1172 + pytest.param(
1173 + _ConvBatchCase(
1174 + spec_query_lens=(1, 4, 4, 0),
1175 + num_accepted=(1, 2, 4, 1),
1176 + spec_query_len=4,
1177 + graph_padding=4,
1178 + ),
1179 + id="spec-graph-padded",
1180 + ),
1181 + pytest.param(
1182 + _ConvBatchCase(
1183 + spec_query_lens=(3, 4),
1184 + num_accepted=(2, 4),
1185 + num_decodes=1,
1186 + prefill_query_lens=(5,),
1187 + spec_query_len=4,
1188 + ),
1189 + id="mixed",
1190 + ),
1191 + pytest.param(
1192 + _ConvBatchCase(
1193 + spec_query_lens=(1, 2),
1194 + num_accepted=(0, 1),
1195 + num_decodes=2,
1196 + prefill_query_lens=(0, 8, 3),
1197 + channels=640,
1198 + kernel_size=3,
1199 + dilation=2,
1200 + spec_query_len=3,
1201 + ),
1202 + id="mixed-varied",
1203 + ),
1204 + ],
1205 +)
1206 +def test_fused_conv_correctness(
1207 + case: _ConvBatchCase,
1208 + state_layout: str,
1209 +) -> None:
1210 + device = torch.device("cuda")
1211 + metadata, num_real_tokens = _make_conv_metadata(case, device)
1212 module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
1213 nn.Module.__init__(module)
404 - module.prefix = "model.layers.1.ple"
405 - inputs = torch.arange(6, dtype=torch.float32).reshape(3, 2)
406 - expected = inputs + 1
407 - monkeypatch.setattr(module, "_short_conv_fallback", lambda _: expected)
408 - monkeypatch.setattr(
409 - ple_layer_module,
410 - "get_forward_context",
411 - lambda: SimpleNamespace(attn_metadata={"model.layers.0.self_attn": object()}),
1214 + module.conv_state_len = (case.kernel_size - 1) * case.dilation
1215 + module.short_conv_dilation = case.dilation
1216 +
1217 + rng, state_reference, conv_state, weights = _make_conv_case(
1218 + device,
1219 + seed=num_real_tokens + case.channels,
1220 + channels=case.channels,
1221 + kernel_size=case.kernel_size,
1222 + dilation=case.dilation,
1223 + state_layout=state_layout,
1224 + spec_query_len=case.spec_query_len,
1225 )
1226 + inputs = torch.randn(
1227 + metadata.num_actual_tokens,
1228 + case.channels,
1229 + device=device,
1230 + dtype=torch.bfloat16,
1231 + generator=rng,
1232 + )
1233 + residual = torch.randn(
1234 + inputs.shape,
1235 + device=device,
1236 + dtype=torch.bfloat16,
1237 + generator=rng,
1238 + )
1239 + null_state = conv_state[NULL_BLOCK_ID].clone()
1240 + residual_kernel = residual.clone()
1241 + residual_reference = residual.clone()
1242 +
1243 + module._short_conv_dilated_dispatch(
1244 + inputs=inputs,
1245 + residual=residual_kernel,
1246 + metadata=metadata,
1247 + conv_state=conv_state,
1248 + conv_weights=weights,
1249 + )
1250 + _short_conv_dilated_dispatch_pytorch(
1251 + inputs=inputs,
1252 + residual=residual_reference,
1253 + metadata=metadata,
1254 + conv_state=state_reference,
1255 + conv_weights=weights,
1256 + conv_state_len=module.conv_state_len,
1257 + dilation=module.short_conv_dilation,
1258 + )
1259 +
1260 + torch.testing.assert_close(
1261 + residual_kernel.float(), residual_reference.float(), atol=3e-2, rtol=3e-2
1262 + )
1263 + assert torch.equal(conv_state, state_reference)
1264 + assert torch.equal(conv_state[NULL_BLOCK_ID], null_state)
1265 + if case.graph_padding:
1266 + assert torch.equal(
1267 + residual_kernel[num_real_tokens:], residual[num_real_tokens:]
1268 + )
1269 +
1270
414 - output = module._short_conv(inputs)
1271 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused gate needs CUDA")
1272 +@pytest.mark.parametrize(("num_tokens", "strided_kv"), [(1, False), (64, True)])
1273 +def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None:
1274 + import math
1275
416 - assert output is expected
1276 + from vllm.models.qwen4_exp.nvidia.ops.ple import ple_gate
1277 +
1278 + device = torch.device("cuda")
1279 + hc, h = 4, 2560
1280 + generator = torch.Generator(device=device).manual_seed(num_tokens)
1281 + kv = torch.randn(
1282 + num_tokens,
1283 + hc * h + h,
1284 + device=device,
1285 + dtype=torch.bfloat16,
1286 + generator=generator,
1287 + )
1288 + key = kv[:, : hc * h]
1289 + value = kv[:, hc * h :]
1290 + if not strided_kv:
1291 + key = key.contiguous()
1292 + value = value.contiguous()
1293 + hidden = torch.randn(
1294 + num_tokens,
1295 + hc * h,
1296 + device=device,
1297 + dtype=torch.bfloat16,
1298 + generator=generator,
1299 + )
1300 +
1301 + def make_norm_weight() -> torch.Tensor:
1302 + weight = torch.empty(hc * h, device=device, dtype=torch.bfloat16)
1303 + return weight.normal_(mean=-0.1, std=0.1, generator=generator)
1304 +
1305 + norm_key = make_norm_weight()
1306 + norm_query = make_norm_weight()
1307 + norm_conv = make_norm_weight()
1308 +
1309 + def grouped_norm(inputs: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
1310 + grouped = inputs.float().unflatten(-1, (hc, h))
1311 + var = grouped.square().mean(dim=-1, keepdim=True)
1312 + normalized = grouped * torch.rsqrt(var + 1e-6)
1313 + return (normalized.flatten(-2) * (1.0 + weight.float())).to(inputs.dtype)
1314 +
1315 + gated, normed = ple_gate(
1316 + key,
1317 + value,
1318 + hidden,
1319 + norm_key,
1320 + norm_query,
1321 + norm_conv,
1322 + 1e-6,
1323 + )
1324 + key_normalized = grouped_norm(key, norm_key).reshape(num_tokens, hc, h)
1325 + query_normalized = grouped_norm(hidden, norm_query).reshape(num_tokens, hc, h)
1326 + dot = (key_normalized * query_normalized).sum(dim=-1, keepdim=True)
1327 + dot = (dot / math.sqrt(h)).to(torch.bfloat16)
1328 + gate = torch.sigmoid(dot.sign() * dot.abs().clamp_min(1e-6).sqrt()).to(
1329 + torch.bfloat16
1330 + )
1331 + expected_gated = (gate * value.unsqueeze(-2)).flatten(-2)
1332 + expected_normed = grouped_norm(expected_gated, norm_conv)
1333 + assert torch.equal(gated, expected_gated)
1334 + torch.testing.assert_close(normed, expected_normed, atol=1e-2, rtol=1e-2)
vllm/models/qwen4_exp/nvidia/model.py
+9 -5
@@ -141,9 +141,9 @@ _QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [
141 "_input_scale",
142 ]
143
144 -# The checkpoint keeps down and injection projections separate; runtime packs
145 -# them into adjacent logical shards of one MergedColumnParallelLinear.
146 -_HC_WEIGHTS_MAPPER = WeightsMapper(
144 +# The checkpoint stores these projections separately; runtime packs each group
145 +# into adjacent logical shards of a MergedColumnParallelLinear.
146 +_EXTRA_WEIGHTS_MAPPER = WeightsMapper(
147 orig_to_new_stacked={
148 "hyper_connection.input_mix_weight_down.weight": (
149 "hyper_connection.input_mix_weight_down_block_inject.weight",
@@ -153,6 +153,8 @@ _HC_WEIGHTS_MAPPER = WeightsMapper(
153 "hyper_connection.input_mix_weight_down_block_inject.weight",
154 1,
155 ),
156 + "ple.key_proj": ("ple.kv_proj", 0),
157 + "ple.value_proj": ("ple.kv_proj", 1),
158 }
159 )
160
@@ -389,7 +391,7 @@ class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
391 }
392 )
393 class Qwen4ExpModel(nn.Module):
392 - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
394 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
395
396 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
397 super().__init__()
@@ -612,6 +614,7 @@ class Qwen4ExpForCausalLM(
614 packed_modules_mapping = {
615 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
616 "gate_up_proj": ["gate_proj", "up_proj"],
617 + "kv_proj": ["key_proj", "value_proj"],
618 "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
619 "in_proj_ba": ["in_proj_b", "in_proj_a"],
620 "input_mix_weight_down_block_inject": [
@@ -848,11 +851,12 @@ class Qwen4ExpForConditionalGeneration(
851 requires_raw_input_tokens = True
852
853 packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
854 + "kv_proj": ["key_proj", "value_proj"],
855 "input_mix_weight_down_block_inject": [
856 "input_mix_weight_down",
857 "block_inject_weight",
858 "_input_mix_padding",
855 - ]
859 + ],
860 }
861
862 @staticmethod
vllm/models/qwen4_exp/nvidia/mtp.py
+2 -2
@@ -52,7 +52,7 @@ from vllm.transformers_utils.configs.qwen4_exp import (
52 from .hyperconnection import GatedResidual, HyperConnectionConfig
53 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
54 from .model import (
55 - _HC_WEIGHTS_MAPPER,
55 + _EXTRA_WEIGHTS_MAPPER,
56 _QWEN4_EXP_IGNORED_MISSING_SUFFIXES,
57 Qwen4ExpDecoderLayer,
58 Qwen4ExpMixtureOfExperts,
@@ -157,7 +157,7 @@ def _make_draft_vllm_config(
157 }
158 )
159 class Qwen4ExpMultiTokenPredictor(nn.Module):
160 - hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
160 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
161
162 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
163 super().__init__()
vllm/models/qwen4_exp/nvidia/ops/ple.py new
+669
@@ -0,0 +1,669 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused Qwen4Exp PLE kernels.
4 +
5 +N-gram IDs hash each suffix with
6 +``offset[h] + (xor_i(token[t-i] * multiplier[i]) % size[h])``. The gate computes
7 +``d = dot(RMSNorm(key), RMSNorm(query)) / sqrt(H)`` and
8 +``g = sigmoid(sign(d) * sqrt(max(abs(d), 1e-6)))``. Short convolution adds
9 +``silu(sum_k weight[k] * history[t + k * dilation])`` to the gated output and
10 +updates its persistent history.
11 +"""
12 +
13 +from typing import Literal
14 +
15 +import torch
16 +
17 +from vllm.platforms import current_platform
18 +from vllm.triton_utils import tl, triton
19 +from vllm.utils.torch_utils import direct_register_custom_op
20 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
21 +
22 +
23 +# N-gram ID generation
24 +@triton.jit(do_not_specialize=["num_tokens", "num_reqs", "binary_search_iters"])
25 +def _ple_ngram_ids_kernel(
26 + input_ids_ptr,
27 + qsl_ptr,
28 + ctx_ptr,
29 + multipliers_ptr,
30 + sizes_ptr,
31 + offsets_ptr,
32 + out_ptr,
33 + num_tokens,
34 + num_reqs,
35 + eos_token_id,
36 + binary_search_iters,
37 + NGRAM_CONTEXT_LEN: tl.constexpr,
38 + HEADS_PER_NGRAM: tl.constexpr,
39 + BLOCK_T: tl.constexpr,
40 + launch_pdl: tl.constexpr,
41 +):
42 + NGRAM_HEADS: tl.constexpr = NGRAM_CONTEXT_LEN * HEADS_PER_NGRAM
43 + BLOCK_H: tl.constexpr = triton.next_power_of_2(NGRAM_HEADS)
44 + pid = tl.program_id(0)
45 + token_offsets = pid * BLOCK_T + tl.arange(0, BLOCK_T)
46 + token_mask = token_offsets < num_tokens
47 +
48 + if launch_pdl:
49 + tl.extra.cuda.gdc_wait()
50 +
51 + # A flattened token offset does not identify its request when query chunks
52 + # have different lengths. Binary-search the request boundaries instead.
53 + p_lo = tl.full([BLOCK_T], 1, tl.int32)
54 + p_hi = tl.full([BLOCK_T], num_reqs + 1, tl.int32)
55 + for _ in range(binary_search_iters):
56 + mid = (p_lo + p_hi) // 2
57 + qmid = tl.load(qsl_ptr + mid, mask=token_mask & (mid <= num_reqs), other=0)
58 + pred = qmid <= token_offsets
59 + p_lo = tl.where(pred, mid + 1, p_lo)
60 + p_hi = tl.where(pred, p_hi, mid)
61 + req = tl.minimum(p_lo - 1, num_reqs - 1).to(tl.int64)
62 + request_start = tl.load(qsl_ptr + req, mask=token_mask, other=0)
63 + chunk_pos = token_offsets - request_start
64 +
65 + # Initialize every head with the current-token term.
66 + current_token = tl.load(input_ids_ptr + token_offsets, mask=token_mask, other=0).to(
67 + tl.int64
68 + )
69 + current_multiplier = tl.load(multipliers_ptr)
70 + mixed = current_token[:, None] * current_multiplier
71 +
72 + g = tl.arange(0, BLOCK_H)
73 + head_mask = g < NGRAM_HEADS
74 + ngram_order = g // HEADS_PER_NGRAM + 2
75 +
76 + # Walk predecessors from newest to oldest. At chunk boundaries they come
77 + # from ngram_context; otherwise they come from this step's input_ids.
78 + crossed = tl.zeros([BLOCK_T], tl.int1)
79 + for shift in tl.static_range(1, NGRAM_CONTEXT_LEN + 1):
80 + in_step = chunk_pos >= shift
81 + ctx_col = NGRAM_CONTEXT_LEN - shift + chunk_pos
82 + step_token = tl.load(
83 + input_ids_ptr + token_offsets - shift,
84 + mask=token_mask & in_step,
85 + other=0,
86 + )
87 + context_token = tl.load(
88 + ctx_ptr + req * NGRAM_CONTEXT_LEN + ctx_col,
89 + mask=token_mask & (~in_step),
90 + other=0,
91 + )
92 + candidate = tl.where(in_step, step_token, context_token).to(tl.int64)
93 + # Older positions remain behind the first EOS boundary.
94 + candidate = tl.where(crossed, eos_token_id, candidate)
95 + crossed = crossed | (candidate == eos_token_id)
96 + multiplier = tl.load(multipliers_ptr + shift)
97 + term = candidate[:, None] * multiplier
98 + mixed = mixed ^ tl.where((ngram_order > shift)[None, :], term, 0)
99 +
100 + # Map each hash into its embedding-table partition.
101 + sizes = tl.load(sizes_ptr + g, mask=head_mask, other=1)[None, :]
102 + head_offsets = tl.load(offsets_ptr + g, mask=head_mask, other=0)[None, :]
103 + # Hash products may overflow int64; preserve torch.remainder semantics.
104 + remainders = mixed % sizes
105 + remainders = tl.where(remainders < 0, remainders + sizes, remainders)
106 + ids = remainders + head_offsets
107 + if launch_pdl:
108 + tl.extra.cuda.gdc_launch_dependents()
109 + tl.store(
110 + out_ptr + token_offsets[:, None] * NGRAM_HEADS + g[None, :],
111 + ids,
112 + mask=token_mask[:, None] & head_mask[None, :],
113 + )
114 +
115 +
116 +def _ple_ngram_ids(
117 + input_ids: torch.Tensor,
118 + query_start_loc: torch.Tensor,
119 + ngram_context: torch.Tensor,
120 + layer_multipliers: torch.Tensor,
121 + ngram_heads_vocab_sizes: torch.Tensor,
122 + ngram_heads_offsets: torch.Tensor,
123 + output: torch.Tensor,
124 + eos_token_id: int,
125 + heads_per_ngram: int,
126 +) -> None:
127 + input_ids = input_ids.reshape(-1)
128 + num_tokens = input_ids.shape[0]
129 + num_reqs = query_start_loc.numel() - 1
130 + ctx_len = ngram_context.shape[1]
131 + BLOCK_T = 8
132 + launch_pdl = current_platform.is_arch_support_pdl()
133 + _ple_ngram_ids_kernel[(triton.cdiv(num_tokens, BLOCK_T),)](
134 + input_ids,
135 + query_start_loc,
136 + ngram_context,
137 + layer_multipliers,
138 + ngram_heads_vocab_sizes,
139 + ngram_heads_offsets,
140 + output,
141 + num_tokens,
142 + num_reqs,
143 + eos_token_id,
144 + binary_search_iters=num_reqs.bit_length(),
145 + NGRAM_CONTEXT_LEN=ctx_len,
146 + HEADS_PER_NGRAM=heads_per_ngram,
147 + BLOCK_T=BLOCK_T,
148 + launch_pdl=launch_pdl,
149 + num_warps=4,
150 + )
151 +
152 +
153 +def ple_ngram_ids(
154 + input_ids: torch.Tensor,
155 + query_start_loc: torch.Tensor,
156 + ngram_context: torch.Tensor,
157 + layer_multipliers: torch.Tensor,
158 + ngram_heads_vocab_sizes: torch.Tensor,
159 + ngram_heads_offsets: torch.Tensor,
160 + eos_token_id: int,
161 + heads_per_ngram: int,
162 + output: torch.Tensor | None = None,
163 +) -> torch.Tensor:
164 + if output is None:
165 + output = torch.empty(
166 + (input_ids.numel(), ngram_context.shape[1] * heads_per_ngram),
167 + dtype=torch.int64,
168 + device=input_ids.device,
169 + )
170 + _ple_ngram_ids(
171 + input_ids,
172 + query_start_loc,
173 + ngram_context,
174 + layer_multipliers,
175 + ngram_heads_vocab_sizes,
176 + ngram_heads_offsets,
177 + output,
178 + eos_token_id,
179 + heads_per_ngram,
180 + )
181 + return output
182 +
183 +
184 +# Gating and normalization
185 +@triton.jit
186 +def _ple_gate_kernel(
187 + key_ptr, # [T, HC*H] with row stride key_rs
188 + value_ptr, # [T, H] with row stride value_rs
189 + hidden_ptr, # [T, HC*H]
190 + nk_ptr, # [HC*H] norm_key weight
191 + nq_ptr, # [HC*H] norm_query weight
192 + ncw_ptr, # [HC*H] norm_conv weight
193 + gated_ptr, # [T, HC*H] out
194 + normed_ptr, # [T, HC*H] out
195 + key_rs,
196 + value_rs,
197 + eps,
198 + H: tl.constexpr,
199 + HC: tl.constexpr,
200 + BLOCK_H: tl.constexpr,
201 + launch_pdl: tl.constexpr,
202 +):
203 + t = tl.program_id(0)
204 + s = tl.program_id(1)
205 + lanes = tl.arange(0, BLOCK_H)
206 + mask = lanes < H
207 + offs = s * H + lanes
208 + dtype: tl.constexpr = key_ptr.dtype.element_ty
209 +
210 + if launch_pdl:
211 + tl.extra.cuda.gdc_wait()
212 +
213 + k = tl.load(key_ptr + t * key_rs + offs, mask=mask, other=0.0).to(tl.float32)
214 + q = tl.load(hidden_ptr + t * HC * H + offs, mask=mask, other=0.0).to(tl.float32)
215 + nk = tl.load(nk_ptr + offs, mask=mask, other=0.0).to(tl.float32)
216 + nq = tl.load(nq_ptr + offs, mask=mask, other=0.0).to(tl.float32)
217 +
218 + # Match eager materialization at each intermediate tensor boundary.
219 + k_n = (k * tl.rsqrt(tl.sum(k * k) / H + eps) * (1.0 + nk)).to(dtype)
220 + q_n = (q * tl.rsqrt(tl.sum(q * q) / H + eps) * (1.0 + nq)).to(dtype)
221 + products = (k_n.to(tl.float32) * q_n.to(tl.float32)).to(dtype)
222 + dot = tl.sum(products.to(tl.float32)).to(dtype).to(tl.float32)
223 + d = dot / tl.sqrt(float(H))
224 + d = d.to(dtype).to(tl.float32)
225 + sign = tl.where(d < 0, -1.0, 0.0)
226 + sign = tl.where(d > 0, 1.0, sign)
227 + magnitude = tl.sqrt(tl.maximum(tl.abs(d), 1e-6)).to(dtype)
228 + g = tl.sigmoid(sign * magnitude)
229 + g = g.to(dtype).to(tl.float32)
230 +
231 + v = tl.load(value_ptr + t * value_rs + lanes, mask=mask, other=0.0).to(tl.float32)
232 + gated = (g * v).to(dtype)
233 + gf = gated.to(tl.float32)
234 + ncw = tl.load(ncw_ptr + offs, mask=mask, other=0.0).to(tl.float32)
235 + normed = gf * tl.rsqrt(tl.sum(gf * gf) / H + eps) * (1.0 + ncw)
236 +
237 + if launch_pdl:
238 + tl.extra.cuda.gdc_launch_dependents()
239 + tl.store(gated_ptr + t * HC * H + offs, gated, mask=mask)
240 + tl.store(normed_ptr + t * HC * H + offs, normed, mask=mask)
241 +
242 +
243 +def _ple_gate(
244 + key: torch.Tensor,
245 + value: torch.Tensor,
246 + hidden: torch.Tensor,
247 + norm_key_w: torch.Tensor,
248 + norm_query_w: torch.Tensor,
249 + norm_conv_w: torch.Tensor,
250 + eps: float,
251 +) -> tuple[torch.Tensor, torch.Tensor]:
252 + num_tokens = hidden.shape[0]
253 + h = value.shape[-1]
254 + hc = hidden.shape[-1] // h
255 + assert key.stride(1) == 1 and value.stride(1) == 1
256 + assert hidden.is_contiguous()
257 + if key.dtype != value.dtype or key.dtype != hidden.dtype:
258 + raise ValueError("key, value, and hidden must have the same dtype")
259 + if key.dtype not in (torch.bfloat16, torch.float16):
260 + raise ValueError("PLE gate supports BF16 and FP16 inputs")
261 + gated = torch.empty_like(hidden)
262 + normed = torch.empty_like(hidden)
263 + _ple_gate_kernel[(num_tokens, hc)](
264 + key,
265 + value,
266 + hidden,
267 + norm_key_w,
268 + norm_query_w,
269 + norm_conv_w,
270 + gated,
271 + normed,
272 + key.stride(0),
273 + value.stride(0),
274 + eps,
275 + H=h,
276 + HC=hc,
277 + BLOCK_H=triton.next_power_of_2(h),
278 + num_warps=4,
279 + launch_pdl=current_platform.is_arch_support_pdl(),
280 + )
281 + return gated, normed
282 +
283 +
284 +def _ple_gate_fake(
285 + key: torch.Tensor,
286 + value: torch.Tensor,
287 + hidden: torch.Tensor,
288 + norm_key_w: torch.Tensor,
289 + norm_query_w: torch.Tensor,
290 + norm_conv_w: torch.Tensor,
291 + eps: float,
292 +) -> tuple[torch.Tensor, torch.Tensor]:
293 + return torch.empty_like(hidden), torch.empty_like(hidden)
294 +
295 +
296 +direct_register_custom_op(
297 + op_name="qwen4_exp_ple_gate",
298 + op_func=_ple_gate,
299 + mutates_args=[],
300 + fake_impl=_ple_gate_fake,
301 +)
302 +
303 +
304 +def ple_gate(
305 + key: torch.Tensor,
306 + value: torch.Tensor,
307 + hidden: torch.Tensor,
308 + norm_key_w: torch.Tensor,
309 + norm_query_w: torch.Tensor,
310 + norm_conv_w: torch.Tensor,
311 + eps: float,
312 +) -> tuple[torch.Tensor, torch.Tensor]:
313 + return torch.ops.vllm.qwen4_exp_ple_gate(
314 + key, value, hidden, norm_key_w, norm_query_w, norm_conv_w, eps
315 + )
316 +
317 +
318 +# Dilated short convolution
319 +@triton.jit(do_not_specialize=["num_reqs", "bs_iters", "has_token_map"])
320 +def _ple_conv_kernel(
321 + x_ptr,
322 + state_ptr,
323 + w_ptr,
324 + residual_ptr,
325 + state_idx_ptr,
326 + qsl_ptr,
327 + num_acc_ptr,
328 + has_init_ptr,
329 + token_idx_ptr,
330 + has_token_map,
331 + num_reqs,
332 + bs_iters,
333 + state_bs,
334 + state_ws,
335 + state_cs,
336 + C: tl.constexpr,
337 + BLOCK_C: tl.constexpr,
338 + STATE_LEN: tl.constexpr,
339 + DILATION: tl.constexpr,
340 + KERNEL_SIZE: tl.constexpr,
341 + SPEC_QUERY_LEN: tl.constexpr,
342 + MODE: tl.constexpr,
343 + HAS_INIT: tl.constexpr,
344 + NULL_STATE_ID: tl.constexpr,
345 + launch_pdl: tl.constexpr,
346 +):
347 + t = tl.program_id(0)
348 + pid_c = tl.program_id(1)
349 + c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
350 + c_mask = c_offs < C
351 +
352 + if launch_pdl:
353 + tl.extra.cuda.gdc_wait()
354 +
355 + # Mode-specific launches use logical row indices. Mixed batches map them
356 + # back to the original token order without gathering inputs or residuals.
357 + output_t = tl.load(token_idx_ptr + t, mask=has_token_map, other=t).to(tl.int64)
358 +
359 + if MODE == "decode":
360 + r = t
361 + q_start = t
362 + j = tl.full([], 0, tl.int32)
363 + slot_off = tl.full([], 0, tl.int32)
364 + in_range = True
365 + else:
366 + # Locate the request containing this token.
367 + lo = tl.full([], 1, tl.int32)
368 + hi = tl.full([], num_reqs + 1, tl.int32)
369 + for _ in range(bs_iters):
370 + mid = (lo + hi) // 2
371 + qmid = tl.load(qsl_ptr + mid, mask=mid <= num_reqs, other=0)
372 + pred = qmid <= t
373 + lo = tl.where(pred, mid + 1, lo)
374 + hi = tl.where(pred, hi, mid)
375 + r = tl.minimum(lo - 1, num_reqs - 1)
376 + q_start = tl.load(qsl_ptr + r)
377 + j = (t - q_start).to(tl.int32)
378 + total_real = tl.load(qsl_ptr + num_reqs)
379 + in_range = t < total_real
380 + if MODE == "spec":
381 + num_acc = tl.load(num_acc_ptr + r)
382 + # Roll back draft state to the last accepted token.
383 + slot_off = tl.minimum(tl.maximum(num_acc - 1, 0), SPEC_QUERY_LEN - 1).to(
384 + tl.int32
385 + )
386 + else:
387 + slot_off = tl.full([], 0, tl.int32)
388 +
389 + sid = tl.load(state_idx_ptr + r).to(tl.int64)
390 + state_ok = sid != NULL_STATE_ID
391 + sid_safe = tl.where(state_ok, sid, 0)
392 + if HAS_INIT:
393 + has_init = tl.load(has_init_ptr + r, mask=state_ok, other=0) != 0
394 + else:
395 + has_init = state_ok
396 + if MODE == "spec":
397 + # A new spec chunk can produce output from its own tokens without a
398 + # prior state; null slots in other modes denote inactive rows.
399 + read_state = state_ok
400 + out_ok = in_range
401 + else:
402 + read_state = state_ok & has_init
403 + out_ok = in_range & state_ok
404 +
405 + # state[sid, w, c] may be a strided view into a shared pool.
406 + base_state = state_ptr + sid_safe * state_bs
407 + acc = tl.zeros([BLOCK_C], tl.float32)
408 + for k in tl.static_range(0, KERNEL_SIZE):
409 + h = j + DILATION * k
410 + from_state = h <= STATE_LEN - 1
411 + state_tap = tl.load(
412 + base_state + (h + slot_off) * state_ws + c_offs * state_cs,
413 + mask=c_mask & read_state & from_state,
414 + other=0.0,
415 + )
416 + input_t = q_start + h - STATE_LEN
417 + input_t = tl.load(
418 + token_idx_ptr + input_t,
419 + mask=has_token_map & out_ok & (~from_state),
420 + other=input_t,
421 + ).to(tl.int64)
422 + input_tap = tl.load(
423 + x_ptr + input_t * C + c_offs,
424 + mask=c_mask & out_ok & (~from_state),
425 + other=0.0,
426 + )
427 + tap = tl.where(from_state, state_tap, input_tap).to(tl.float32)
428 + weight = tl.load(
429 + w_ptr + c_offs * KERNEL_SIZE + k,
430 + mask=c_mask,
431 + other=0.0,
432 + ).to(tl.float32)
433 + acc += weight * tap
434 +
435 + # F.conv1d materializes its output dtype before SiLU.
436 + conv = acc.to(residual_ptr.dtype.element_ty).to(tl.float32)
437 + y = conv * tl.sigmoid(conv)
438 + conv_output = tl.where(out_ok, y, 0.0).to(residual_ptr.dtype.element_ty)
439 + residual = tl.load(
440 + residual_ptr + output_t * C + c_offs,
441 + mask=c_mask,
442 + other=0.0,
443 + )
444 + if launch_pdl:
445 + tl.extra.cuda.gdc_launch_dependents()
446 + tl.store(
447 + residual_ptr + output_t * C + c_offs,
448 + residual + conv_output,
449 + mask=c_mask,
450 + )
451 +
452 + if MODE == "decode":
453 + # Decode has one program per request and channel block, so no peer can
454 + # still be reading this state slice when it is updated.
455 + decode_input = tl.load(
456 + x_ptr + output_t * C + c_offs,
457 + mask=c_mask & state_ok,
458 + other=0.0,
459 + )
460 + for i in tl.static_range(0, STATE_LEN):
461 + if i < STATE_LEN - 1:
462 + next_state = tl.load(
463 + base_state + (i + 1) * state_ws + c_offs * state_cs,
464 + mask=c_mask & state_ok & has_init,
465 + other=0.0,
466 + )
467 + else:
468 + next_state = decode_input
469 + tl.store(
470 + base_state + i * state_ws + c_offs * state_cs,
471 + next_state,
472 + mask=c_mask & state_ok,
473 + )
474 +
475 +
476 +# Prefill and spec write back separately so every token reads the old state.
477 +@triton.jit(do_not_specialize=["has_token_map"])
478 +def _ple_conv_writeback_kernel(
479 + x_ptr,
480 + state_ptr,
481 + state_idx_ptr,
482 + qsl_ptr,
483 + num_acc_ptr,
484 + has_init_ptr,
485 + token_idx_ptr,
486 + has_token_map,
487 + state_bs,
488 + state_ws,
489 + state_cs,
490 + C: tl.constexpr,
491 + BLOCK_C: tl.constexpr,
492 + STATE_LEN: tl.constexpr,
493 + SPEC_QUERY_LEN: tl.constexpr,
494 + STATE_WIDTH: tl.constexpr,
495 + MODE: tl.constexpr,
496 + HAS_INIT: tl.constexpr,
497 + NULL_STATE_ID: tl.constexpr,
498 +):
499 + r = tl.program_id(0)
500 + pid_c = tl.program_id(1)
501 + c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
502 + c_mask = c_offs < C
503 +
504 + sid = tl.load(state_idx_ptr + r).to(tl.int64)
505 + state_ok = sid != NULL_STATE_ID
506 + if not state_ok:
507 + return
508 +
509 + q_start = tl.load(qsl_ptr + r)
510 + q_end = tl.load(qsl_ptr + r + 1)
511 + qlen = (q_end - q_start).to(tl.int32)
512 + if MODE == "spec":
513 + num_acc = tl.load(num_acc_ptr + r)
514 + # Use the same rollback window as the output kernel.
515 + slot_off = tl.minimum(tl.maximum(num_acc - 1, 0), SPEC_QUERY_LEN - 1).to(
516 + tl.int32
517 + )
518 + shift = 1
519 + else:
520 + slot_off = tl.full([], 0, tl.int32)
521 + shift = qlen
522 + if qlen <= 0:
523 + return
524 +
525 + has_init = tl.load(has_init_ptr + r) != 0 if HAS_INIT else True
526 + src_ok = state_ok if MODE == "spec" else state_ok & has_init
527 +
528 + WRITE_W: tl.constexpr = STATE_WIDTH if MODE == "spec" else STATE_LEN
529 + base_state = state_ptr + sid * state_bs
530 + for i in tl.static_range(0, WRITE_W):
531 + m = shift + i
532 + from_state = m <= STATE_LEN - 1
533 + # Short or graph-padded spec chunks preserve the unused state tail.
534 + do_write = i < STATE_LEN + qlen - 1 if MODE == "spec" else True
535 + state_value = tl.load(
536 + base_state + (slot_off + m) * state_ws + c_offs * state_cs,
537 + mask=c_mask & from_state & src_ok,
538 + other=0.0,
539 + )
540 + input_t = q_start + m - STATE_LEN
541 + input_t = tl.load(
542 + token_idx_ptr + input_t,
543 + mask=has_token_map & (~from_state) & do_write,
544 + other=input_t,
545 + ).to(tl.int64)
546 + input_value = tl.load(
547 + x_ptr + input_t * C + c_offs,
548 + mask=c_mask & (~from_state) & do_write,
549 + other=0.0,
550 + )
551 + value = tl.where(from_state, state_value, input_value)
552 + tl.store(
553 + base_state + i * state_ws + c_offs * state_cs,
554 + value,
555 + mask=c_mask & do_write,
556 + )
557 +
558 +
559 +def ple_conv(
560 + inputs: torch.Tensor,
561 + residual: torch.Tensor,
562 + conv_state: torch.Tensor,
563 + conv_weights: torch.Tensor,
564 + state_indices: torch.Tensor,
565 + *,
566 + mode: Literal["decode", "spec", "prefill"],
567 + dilation: int,
568 + query_start_loc: torch.Tensor | None = None,
569 + num_accepted_tokens: torch.Tensor | None = None,
570 + has_initial_states: torch.Tensor | None = None,
571 + spec_query_len: int = 1,
572 + token_indices: torch.Tensor | None = None,
573 +) -> None:
574 + """Add short-convolution output to ``residual`` and update its state."""
575 + BLOCK_C = 512
576 + kernel_spec_query_len = spec_query_len if mode == "spec" else 1
577 + T, C = inputs.shape
578 + K = conv_weights.shape[1]
579 + state_len = (K - 1) * dilation
580 + state_width = state_len + kernel_spec_query_len - 1
581 + if token_indices is not None:
582 + T = token_indices.numel()
583 + if conv_state.shape[1] != C or conv_state.shape[2] < state_width:
584 + raise ValueError(
585 + "conv_state must have shape [slots, channels, window], with "
586 + f"channels={C} and window >= {state_width}"
587 + )
588 + state_bs, state_cs, state_ws = conv_state.stride()
589 +
590 + if mode == "decode":
591 + num_reqs = T
592 + binary_search_iters = 1
593 + has_initial_states_arg = has_initial_states is not None
594 + elif mode == "spec":
595 + if query_start_loc is None or num_accepted_tokens is None:
596 + raise ValueError(
597 + "query_start_loc and num_accepted_tokens are required for spec decode"
598 + )
599 + num_reqs = state_indices.numel()
600 + binary_search_iters = max(num_reqs, 1).bit_length()
601 + has_initial_states_arg = False
602 + elif mode == "prefill":
603 + if query_start_loc is None or has_initial_states is None:
604 + raise ValueError(
605 + "query_start_loc and has_initial_states are required for prefill"
606 + )
607 + num_reqs = state_indices.numel()
608 + binary_search_iters = max(num_reqs, 1).bit_length()
609 + has_initial_states_arg = True
610 + else:
611 + raise ValueError(f"Unsupported short-conv mode: {mode}")
612 +
613 + num_warps = 4 if mode == "prefill" else 8
614 + launch_pdl = current_platform.is_arch_support_pdl()
615 +
616 + # Constexpr flags eliminate accesses to optional None arguments. Without a
617 + # token map, state_indices is an unused but device-resident placeholder.
618 + _ple_conv_kernel[(T, triton.cdiv(C, BLOCK_C))](
619 + inputs,
620 + conv_state,
621 + conv_weights,
622 + residual,
623 + state_indices,
624 + query_start_loc,
625 + num_accepted_tokens,
626 + has_initial_states,
627 + token_indices if token_indices is not None else state_indices,
628 + token_indices is not None,
629 + num_reqs,
630 + binary_search_iters,
631 + state_bs,
632 + state_ws,
633 + state_cs,
634 + C=C,
635 + BLOCK_C=BLOCK_C,
636 + STATE_LEN=state_len,
637 + DILATION=dilation,
638 + KERNEL_SIZE=K,
639 + SPEC_QUERY_LEN=kernel_spec_query_len,
640 + MODE=mode,
641 + HAS_INIT=has_initial_states_arg,
642 + NULL_STATE_ID=NULL_BLOCK_ID,
643 + launch_pdl=launch_pdl,
644 + num_warps=num_warps,
645 + )
646 + # conv state update is fused with the kernel above for decode
647 + if mode != "decode":
648 + _ple_conv_writeback_kernel[(num_reqs, triton.cdiv(C, BLOCK_C))](
649 + inputs,
650 + conv_state,
651 + state_indices,
652 + query_start_loc,
653 + num_accepted_tokens,
654 + has_initial_states,
655 + token_indices if token_indices is not None else state_indices,
656 + token_indices is not None,
657 + state_bs,
658 + state_ws,
659 + state_cs,
660 + C=C,
661 + BLOCK_C=BLOCK_C,
662 + STATE_LEN=state_len,
663 + SPEC_QUERY_LEN=kernel_spec_query_len,
664 + STATE_WIDTH=state_width,
665 + MODE=mode,
666 + HAS_INIT=has_initial_states_arg,
667 + NULL_STATE_ID=NULL_BLOCK_ID,
668 + num_warps=num_warps,
669 + )
vllm/models/qwen4_exp/nvidia/ple_layer.py
+173 -516
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """GPU-resident Qwen4Exp position-learning enhancement layers."""
4
5 -import math
5 from collections.abc import Iterable, Sequence
6
7 import torch
@@ -11,7 +10,7 @@ from torch import nn
10
11 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
12 from vllm.forward_context import get_forward_context
14 -from vllm.model_executor.layers.linear import ReplicatedLinear
13 +from vllm.model_executor.layers.linear import MergedColumnParallelLinear
14 from vllm.model_executor.layers.mamba.abstract import MambaBase
15 from vllm.model_executor.layers.mamba.mamba_utils import (
16 MambaStateDtypeCalculator,
@@ -42,9 +41,9 @@ from vllm.v1.attention.backends.short_conv_attn import (
41 PleShortConvAttentionBackend,
42 PleShortConvAttentionMetadata,
43 )
45 -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
44
45 from ..common.ple import PLEVocabParallelEmbedding
46 +from .ops.ple import ple_conv, ple_gate, ple_ngram_ids
47
48
49 class Qwen4ExpPLEGroupedNorm(nn.Module):
@@ -256,8 +255,6 @@ class Qwen4ExpNGramEmbedding(nn.Module):
255 config: Qwen4ExpTextConfig,
256 embedding_dim: int,
257 ple_dense_layer_id: int,
259 - max_total_tokens: int,
260 - max_num_reqs: int,
258 prefix: str,
259 layer_name: str,
260 quant_config: QuantizationConfig | None = None,
@@ -324,20 +321,6 @@ class Qwen4ExpNGramEmbedding(nn.Module):
321 quant_config, f"{prefix}.ngram_embedding"
322 ),
323 )
327 - self.register_buffer(
328 - "positions_buffer",
329 - torch.arange(max_total_tokens, dtype=torch.int64),
330 - persistent=False,
331 - )
332 - self.register_buffer(
333 - "padded_buffer",
334 - torch.full(
335 - (max_num_reqs, max_total_tokens),
336 - self.eos_token_id,
337 - dtype=torch.int64,
338 - ),
339 - persistent=False,
340 - )
324
325 @staticmethod
326 def _shift_precompute(
@@ -379,26 +362,34 @@ class Qwen4ExpNGramEmbedding(nn.Module):
362 input_ids: torch.Tensor,
363 query_start_loc: torch.Tensor,
364 ngram_context: torch.Tensor,
365 + output: torch.Tensor | None = None,
366 ) -> torch.Tensor:
367 """Compute n-gram embedding indices for the current request layout."""
384 - input_ids = input_ids.reshape(-1).long()
385 - query_start_loc = query_start_loc.long()
368 + input_ids = input_ids.reshape(-1)
369 num_reqs = query_start_loc.numel() - 1
370 num_tokens = input_ids.shape[0]
388 - if num_tokens > self.positions_buffer.numel():
389 - raise ValueError(
390 - f"PLE received {num_tokens} tokens, but its workspace supports "
391 - f"at most {self.positions_buffer.numel()}"
392 - )
393 - if num_reqs > self.padded_buffer.shape[0]:
394 - raise ValueError(
395 - f"PLE received {num_reqs} requests, but its workspace supports "
396 - f"at most {self.padded_buffer.shape[0]}"
397 - )
371
399 - positions = self.positions_buffer[:num_tokens]
400 - packed = self.padded_buffer[:num_reqs, :num_tokens]
401 - packed.fill_(self.eos_token_id)
372 + if input_ids.is_cuda:
373 + return ple_ngram_ids(
374 + input_ids=input_ids,
375 + query_start_loc=query_start_loc,
376 + ngram_context=ngram_context,
377 + layer_multipliers=self.layer_multipliers,
378 + ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes,
379 + ngram_heads_offsets=self.ngram_heads_offsets,
380 + eos_token_id=self.eos_token_id,
381 + heads_per_ngram=self.heads_per_ngram,
382 + output=output,
383 + )
384 + input_ids = input_ids.long()
385 + query_start_loc = query_start_loc.long()
386 + positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64)
387 + packed = torch.full(
388 + (num_reqs, num_tokens),
389 + self.eos_token_id,
390 + device=input_ids.device,
391 + dtype=torch.int64,
392 + )
393 request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
394 request_indices.clamp_(max=num_reqs - 1)
395 columns = (positions - query_start_loc[request_indices]).clamp(
@@ -446,12 +437,13 @@ class Qwen4ExpNGramEmbedding(nn.Module):
437 query_start_loc: torch.Tensor,
438 ngram_context: torch.Tensor,
439 ) -> torch.Tensor:
449 - ngram_ids = input_ids.new_empty(
450 - (input_ids.shape[0], self.ngram_heads),
451 - dtype=torch.long,
452 - )
440 # Keep num_reqs-dependent ID generation outside PIECEWISE CUDA graphs,
441 # which dispatch only on the padded token count.
442 + # torch.compile requires the splitting op to write graph-owned storage.
443 + # Once compilation is removed, the op can return the IDs directly.
444 + ngram_ids = input_ids.new_empty(
445 + (input_ids.numel(), self.ngram_heads), dtype=torch.long
446 + )
447 torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(
448 input_ids,
449 query_start_loc,
@@ -562,26 +554,20 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
554 config,
555 int(config.ple_embed_dim),
556 self.ple_dense_layer_id,
565 - vllm_config.scheduler_config.max_num_batched_tokens,
566 - vllm_config.scheduler_config.max_num_seqs,
557 prefix=f"{prefix}.ple_embedding",
558 layer_name=prefix,
559 quant_config=quant_config,
560 params_dtype=model_config.dtype,
561 )
572 - self.key_proj = ReplicatedLinear(
562 + # The PLE cache is TP-replicated, so this merged projection is too.
563 + self.kv_proj = MergedColumnParallelLinear(
564 int(config.ple_embed_dim),
574 - self.hc_hidden_size,
575 - bias=False,
576 - quant_config=quant_config,
577 - prefix=f"{prefix}.key_proj",
578 - )
579 - self.value_proj = ReplicatedLinear(
580 - int(config.ple_embed_dim),
581 - self.hidden_size,
565 + [self.hc_hidden_size, self.hidden_size],
566 bias=False,
567 + params_dtype=model_config.dtype,
568 quant_config=quant_config,
584 - prefix=f"{prefix}.value_proj",
569 + prefix=f"{prefix}.kv_proj",
570 + disable_tp=True,
571 )
572 norm_args = (
573 self.hc_hidden_size,
@@ -654,360 +640,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
640 num_spec=self.num_spec_tokens,
641 )
642
657 - def _apply_norm(
658 - self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor
659 - ) -> torch.Tensor:
660 - shape = hidden_states.shape
661 - return norm(hidden_states.flatten(-2)).reshape(shape)
662 -
663 - def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor:
664 - # Profiling / CUDA graph capture only; conv state is not updated.
665 - inputs_t = inputs.transpose(0, 1).unsqueeze(0)
666 - output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)]
667 - return F.silu(output).squeeze(0).transpose(0, 1)
668 -
669 - def _short_conv_dilated_decode_batched(
670 - self,
671 - x_d: torch.Tensor,
672 - conv_state: torch.Tensor,
673 - conv_weights: torch.Tensor,
674 - state_indices_tensor_d: torch.Tensor,
675 - has_initial_states_d: torch.Tensor | None,
676 - ) -> torch.Tensor:
677 - state_indices = state_indices_tensor_d.to(
678 - device=conv_state.device, dtype=torch.int64
679 - )
680 - # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
681 - # slot 0 for a safe gather, then zero output and skip write-back.
682 - valid_state = state_indices != NULL_BLOCK_ID
683 - state_indices = torch.where(
684 - valid_state, state_indices, torch.zeros_like(state_indices)
685 - )
686 - if has_initial_states_d is None:
687 - has_initial_state = valid_state
688 - else:
689 - if has_initial_states_d.numel() < state_indices_tensor_d.numel():
690 - raise ValueError(
691 - "has_initial_states_d size mismatch: "
692 - f"got {has_initial_states_d.numel()}, "
693 - f"need >= {state_indices_tensor_d.numel()}."
694 - )
695 - has_initial_state = has_initial_states_d[
696 - : state_indices_tensor_d.numel()
697 - ].to(device=conv_state.device, dtype=torch.bool)
698 - has_initial_state = has_initial_state & valid_state
699 -
700 - cached_state = conv_state.index_select(0, state_indices)
701 - state = cached_state[..., : self.conv_state_len].to(x_d.dtype)
702 - if self.conv_state_len > 0:
703 - initial_state = torch.where(
704 - has_initial_state.view(-1, 1, 1),
705 - state,
706 - torch.zeros_like(state),
707 - )
708 - history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
709 - else:
710 - history = x_d.unsqueeze(-1)
711 -
712 - conv_output = F.conv1d(
713 - history,
714 - conv_weights.unsqueeze(1).contiguous(),
715 - groups=history.size(1),
716 - dilation=self.short_conv_dilation,
717 - ).squeeze(-1)
718 - output = F.silu(conv_output)
719 - output = output * valid_state.view(-1, 1).to(output.dtype)
720 -
721 - if self.conv_state_len > 0:
722 - next_state = history[..., -self.conv_state_len :]
723 - # Padded rows are remapped to the reserved null slot. Preserve its
724 - # existing value while writing the new states for valid rows.
725 - existing_base_state = cached_state[..., : self.conv_state_len]
726 - safe_next_state = torch.where(
727 - valid_state.view(-1, 1, 1),
728 - next_state.to(conv_state.dtype),
729 - existing_base_state,
730 - )
731 - cached_state[..., : self.conv_state_len] = safe_next_state
732 - conv_state.index_copy_(0, state_indices, cached_state)
733 -
734 - return output
735 -
736 - def _short_conv_dilated_prefill_batched(
737 - self,
738 - x_p: torch.Tensor,
739 - metadata: PleShortConvAttentionMetadata,
740 - conv_state: torch.Tensor,
741 - conv_weights: torch.Tensor,
742 - state_indices_tensor_p: torch.Tensor,
743 - num_prefills: int,
744 - num_decode_tokens: int,
745 - num_prefill_tokens: int,
746 - ) -> torch.Tensor:
747 - # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
748 - # requests and equals ``query_start_loc`` when spec-decode is inactive.
749 - non_spec_query_start_loc = metadata.non_spec_query_start_loc
750 - if non_spec_query_start_loc is None:
751 - raise ValueError("query_start_loc is required for prefill short-conv")
752 - query_start_loc_p = (
753 - non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
754 - )
755 - # The metadata builder guarantees that the prefill query offsets start
756 - # at 0 and end at num_prefill_tokens. Avoid reading those values here,
757 - # since doing so would force a device-to-host synchronization.
758 - has_initial_states_p = metadata.has_initial_states_p
759 - if has_initial_states_p is None:
760 - raise ValueError("has_initial_states_p is required for prefill short-conv")
761 -
762 - output = torch.empty_like(x_p)
763 - q_starts = query_start_loc_p.to(torch.int64)
764 - if state_indices_tensor_p.numel() < num_prefills:
765 - raise ValueError(
766 - "state_indices_tensor_p size mismatch: "
767 - f"got {state_indices_tensor_p.numel()}, "
768 - f"need >= {num_prefills}."
769 - )
770 - if has_initial_states_p.numel() < num_prefills:
771 - raise ValueError(
772 - "has_initial_states_p size mismatch: "
773 - f"got {has_initial_states_p.numel()}, "
774 - f"need >= {num_prefills}."
775 - )
776 - if num_prefills == 0 or x_p.numel() == 0:
777 - return output
778 - lengths = q_starts[1:] - q_starts[:-1]
779 - # Use the CPU-computed packing width from the metadata builder instead
780 - # of synchronizing on lengths.max().
781 - max_len = metadata.max_prefill_query_len
782 - if max_len <= 0:
783 - return output
784 -
785 - hidden_size = x_p.shape[1]
786 - positions = torch.arange(
787 - num_prefill_tokens, device=x_p.device, dtype=torch.int64
788 - )
789 - req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
790 - col_indices = positions - q_starts[req_indices]
791 -
792 - packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
793 - packed_tokens[req_indices, col_indices] = x_p
794 - packed_tokens = packed_tokens.transpose(1, 2).contiguous()
795 -
796 - state_indices = state_indices_tensor_p[:num_prefills].to(
797 - device=conv_state.device, dtype=torch.int64
798 - )
799 - valid_state = state_indices != NULL_BLOCK_ID
800 - state_indices = torch.where(
801 - valid_state, state_indices, torch.zeros_like(state_indices)
802 - )
803 - has_initial = has_initial_states_p[:num_prefills].to(
804 - device=conv_state.device, dtype=torch.bool
805 - )
806 - if self.conv_state_len > 0:
807 - if conv_state.shape[0] == 0:
808 - state = conv_state.new_zeros(
809 - (num_prefills, hidden_size, self.conv_state_len),
810 - dtype=x_p.dtype,
811 - )
812 - else:
813 - state = conv_state.index_select(0, state_indices)[
814 - ..., : self.conv_state_len
815 - ].to(x_p.dtype)
816 - use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
817 - initial_state = torch.where(
818 - use_initial_mask,
819 - state,
820 - torch.zeros_like(state),
821 - )
822 - history = torch.cat((initial_state, packed_tokens), dim=-1)
823 - else:
824 - history = packed_tokens
825 -
826 - conv_output = F.conv1d(
827 - history,
828 - conv_weights.unsqueeze(1).contiguous(),
829 - groups=history.size(1),
830 - dilation=self.short_conv_dilation,
831 - )
832 - conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
833 -
834 - token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
835 - valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
836 - valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
837 - num_prefills, 1
838 - )
839 - conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
840 - output.copy_(conv_output[req_indices, col_indices])
841 -
842 - if self.conv_state_len > 0 and conv_state.shape[0] > 0:
843 - state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
844 - num_prefills, 1, 1
845 - )
846 - state_offsets = torch.arange(
847 - self.conv_state_len, device=history.device, dtype=torch.int64
848 - ).view(1, 1, self.conv_state_len)
849 - next_state = history.gather(
850 - dim=2,
851 - index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
852 - )
853 - # Write back without a host synchronization. Valid, non-empty rows
854 - # receive their new state; padding and zero-length rows keep the
855 - # current cache value.
856 - existing_state = conv_state.index_select(0, state_indices)
857 - existing_base_state = existing_state[..., : self.conv_state_len]
858 - update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
859 - safe_next_state = torch.where(
860 - update_mask.view(num_prefills, 1, 1),
861 - next_state.to(conv_state.dtype),
862 - existing_base_state,
863 - )
864 - existing_state[..., : self.conv_state_len] = safe_next_state
865 - conv_state.index_copy_(0, state_indices, existing_state)
866 - return output
867 -
868 - def _short_conv_dilated_spec_batched(
869 - self,
870 - x_spec: torch.Tensor,
871 - conv_state: torch.Tensor,
872 - conv_weights: torch.Tensor,
873 - spec_state_indices_tensor: torch.Tensor,
874 - spec_query_start_loc: torch.Tensor,
875 - num_accepted_tokens: torch.Tensor,
876 - spec_query_len: int,
877 - ) -> torch.Tensor:
878 - """Dilated short-conv for speculative-decode (MTP) requests.
879 -
880 - Each spec request feeds multiple (draft + 1) query tokens. The conv
881 - outputs are computed causally after rolling back the previous draft
882 - state by ``num_accepted_tokens - 1``. The current candidate inputs stay
883 - in the extended cache for the next forward, matching
884 - ``causal_conv1d_update``.
885 -
886 - ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
887 - length and is a Python int, so no host synchronization is needed; this
888 - keeps the path safe for full CUDA-graph capture/replay where the buffers
889 - are padded at the request level.
890 - """
891 - num_reqs = spec_state_indices_tensor.numel()
892 - hidden_size = x_spec.size(-1)
893 - # Use a fixed packing width instead of synchronizing on lengths.max().
894 - max_len = spec_query_len
895 - # Full CUDA graphs can pad these buffers. Only the first num_reqs
896 - # accepted-token counts belong to actual speculative requests.
897 - num_accepted_tokens = num_accepted_tokens[:num_reqs]
898 - q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
899 - # Keep the number of real speculative tokens on the device.
900 - total_real_tokens = q_starts[num_reqs]
901 -
902 - state_indices = spec_state_indices_tensor.to(
903 - device=conv_state.device, dtype=torch.int64
904 - )
905 - valid_state = state_indices != NULL_BLOCK_ID
906 - state_indices = torch.where(
907 - valid_state, state_indices, torch.zeros_like(state_indices)
908 - )
909 - positions = torch.arange(
910 - x_spec.size(0), device=x_spec.device, dtype=torch.int64
911 - )
912 - # Route graph-padded token rows to the discarded dummy request so that
913 - # they cannot overwrite real packed data.
914 - req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
915 - valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
916 - clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
917 - col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
918 - pack_req_indices = torch.where(
919 - valid_tokens,
920 - clamped_req_indices,
921 - torch.full_like(req_indices, num_reqs),
922 - )
923 - pack_col_indices = torch.where(
924 - valid_tokens, col_indices, torch.zeros_like(col_indices)
925 - )
926 -
927 - # The last request row is the dummy sink for graph padding.
928 - packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
929 - packed[pack_req_indices, pack_col_indices] = x_spec
930 - packed = packed.transpose(1, 2).contiguous()
931 -
932 - if self.conv_state_len > 0:
933 - cached_state = conv_state.index_select(0, state_indices)
934 - rollback_offsets = num_accepted_tokens.to(
935 - device=conv_state.device, dtype=torch.int64
936 - ).sub(1)
937 - rollback_offsets = torch.where(
938 - valid_state,
939 - rollback_offsets.clamp_(0, max_len - 1),
940 - torch.zeros_like(rollback_offsets),
941 - )
942 - state_offsets = torch.arange(
943 - self.conv_state_len, device=conv_state.device, dtype=torch.int64
944 - ).view(1, 1, self.conv_state_len)
945 - rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
946 - state = cached_state.gather(
947 - 2, rollback_indices.expand(-1, hidden_size, -1)
948 - ).to(x_spec.dtype)
949 - state = torch.where(
950 - valid_state.view(num_reqs, 1, 1),
951 - state,
952 - torch.zeros_like(state),
953 - )
954 - # Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
955 - dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len))
956 - state_full = torch.cat((state, dummy_state), dim=0)
957 - history = torch.cat((state_full, packed), dim=-1)
958 - else:
959 - history = packed
960 -
961 - conv_output = F.conv1d(
962 - history,
963 - conv_weights.unsqueeze(1).contiguous(),
964 - groups=history.size(1),
965 - dilation=self.short_conv_dilation,
966 - )
967 - conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
968 -
969 - output = conv_output[pack_req_indices, pack_col_indices]
970 - output = output * valid_tokens.view(-1, 1).to(output.dtype)
971 -
972 - # Keep all current candidate inputs in the extended state. On the next
973 - # target forward, ``num_accepted_tokens - 1`` selects the rollback
974 - # window before processing the newly scheduled tokens.
975 - if self.conv_state_len > 0:
976 - state_capacity = self.conv_state_len + max_len - 1
977 - if conv_state.size(-1) < state_capacity:
978 - raise RuntimeError(
979 - "PLE short-conv cache cannot retain speculative tokens: "
980 - f"got {conv_state.size(-1)}, need {state_capacity}."
981 - )
982 - candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
983 - query_lengths = q_starts[1:] - q_starts[:-1]
984 - state_positions = torch.arange(
985 - state_capacity, device=history.device, dtype=torch.int64
986 - ).view(1, 1, state_capacity)
987 - update_lengths = (self.conv_state_len + query_lengths - 1).view(
988 - num_reqs, 1, 1
989 - )
990 - update_mask = valid_state.view(num_reqs, 1, 1) & (
991 - state_positions < update_lengths
992 - )
993 - existing_state = cached_state[..., :state_capacity]
994 - next_state = torch.where(
995 - update_mask,
996 - candidate_state.to(conv_state.dtype),
997 - existing_state,
998 - )
999 - cached_state[..., :state_capacity] = next_state
1000 - conv_state.index_copy_(0, state_indices, cached_state)
1001 -
1002 - return output
1003 -
643 def _short_conv_dilated_dispatch(
644 self,
645 inputs: torch.Tensor,
646 + residual: torch.Tensor,
647 metadata: PleShortConvAttentionMetadata,
648 conv_state: torch.Tensor,
649 conv_weights: torch.Tensor,
1010 - ) -> torch.Tensor:
650 + ) -> None:
651 num_prefills = metadata.num_prefills
652 num_decodes = metadata.num_decodes
653 num_decode_tokens = metadata.num_decode_tokens
@@ -1015,110 +655,129 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
655 has_prefill = num_prefills > 0
656 has_decode = num_decodes > 0
657 has_spec = metadata.spec_sequence_masks is not None
1018 - x = inputs[: metadata.num_actual_tokens]
658 + has_non_spec = has_prefill or has_decode
659 + inputs = inputs[: metadata.num_actual_tokens]
660 + residual = residual[: metadata.num_actual_tokens]
661
1020 - # Split spec / non-spec tokens.
1021 - if has_spec:
1022 - if has_prefill or has_decode:
1023 - assert metadata.spec_token_indx is not None
1024 - assert metadata.non_spec_token_indx is not None
1025 - x_spec = x.index_select(0, metadata.spec_token_indx.long())
1026 - x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long())
1027 - else:
1028 - x_spec = x
1029 - x_non_spec = None
1030 - else:
1031 - x_spec = None
1032 - x_non_spec = x
662 + spec_token_indices = None
663 + non_spec_token_indices = None
664 + if has_spec and has_non_spec:
665 + assert metadata.spec_token_indx is not None
666 + assert metadata.non_spec_token_indx is not None
667 + spec_token_indices = metadata.spec_token_indx
668 + non_spec_token_indices = metadata.non_spec_token_indx
669
1034 - spec_output = None
1035 - # 1. Run the multi-query speculative-decode part.
670 if has_spec:
671 assert metadata.spec_state_indices_tensor is not None
1038 - assert metadata.spec_query_start_loc is not None
1039 - assert metadata.num_accepted_tokens is not None
1040 - spec_output = self._short_conv_dilated_spec_batched(
1041 - x_spec=x_spec,
672 + query_start_loc = metadata.spec_query_start_loc
673 + num_accepted_tokens = metadata.num_accepted_tokens
674 + assert query_start_loc is not None
675 + assert num_accepted_tokens is not None
676 + spec_state_indices = metadata.spec_state_indices_tensor[
677 + : metadata.num_spec_decodes
678 + ]
679 + # Mixed batches stay in their original row order; the kernels map
680 + # logical spec/non-spec rows instead of materializing both groups.
681 + ple_conv(
682 + inputs=inputs,
683 + residual=residual,
684 conv_state=conv_state,
685 conv_weights=conv_weights,
1044 - spec_state_indices_tensor=metadata.spec_state_indices_tensor[
1045 - : metadata.num_spec_decodes
1046 - ],
1047 - spec_query_start_loc=metadata.spec_query_start_loc,
1048 - num_accepted_tokens=metadata.num_accepted_tokens,
686 + state_indices=spec_state_indices,
687 + mode="spec",
688 + dilation=self.short_conv_dilation,
689 + query_start_loc=query_start_loc,
690 + num_accepted_tokens=num_accepted_tokens,
691 spec_query_len=metadata.spec_query_len,
692 + token_indices=spec_token_indices,
693 )
694
1052 - # 2. Run regular decode and prefill requests.
1053 - conv_out_non_spec = None
1054 - state_indices_tensor = metadata.state_indices_tensor
1055 - if x_non_spec is not None:
1056 - assert state_indices_tensor is not None
1057 - if has_prefill:
1058 - state_indices_tensor_d, state_indices_tensor_p = torch.split(
1059 - state_indices_tensor,
1060 - [num_decodes, num_prefills],
1061 - dim=0,
695 + if not has_non_spec:
696 + return
697 +
698 + state_indices = metadata.state_indices_tensor
699 + assert state_indices is not None
700 + if has_prefill:
701 + state_indices_d, state_indices_p = torch.split(
702 + state_indices, [num_decodes, num_prefills], dim=0
703 + )
704 + if non_spec_token_indices is None:
705 + inputs_d, inputs_p = torch.split(
706 + inputs, [num_decode_tokens, num_prefill_tokens], dim=0
707 )
1063 - x_d, x_p = torch.split(
1064 - x_non_spec,
708 + residual_d, residual_p = torch.split(
709 + residual, [num_decode_tokens, num_prefill_tokens], dim=0
710 + )
711 + token_indices_d = None
712 + token_indices_p = None
713 + else:
714 + inputs_d = inputs_p = inputs
715 + residual_d = residual_p = residual
716 + token_indices_d, token_indices_p = torch.split(
717 + non_spec_token_indices,
718 [num_decode_tokens, num_prefill_tokens],
719 dim=0,
720 )
1068 - non_spec_parts: list[torch.Tensor] = []
1069 - if has_decode:
1070 - non_spec_parts.append(
1071 - self._short_conv_dilated_decode_batched(
1072 - x_d=x_d,
1073 - conv_state=conv_state,
1074 - conv_weights=conv_weights,
1075 - state_indices_tensor_d=state_indices_tensor_d,
1076 - has_initial_states_d=metadata.has_initial_states_d,
1077 - )
1078 - )
1079 - non_spec_parts.append(
1080 - self._short_conv_dilated_prefill_batched(
1081 - x_p=x_p,
1082 - metadata=metadata,
1083 - conv_state=conv_state,
1084 - conv_weights=conv_weights,
1085 - state_indices_tensor_p=state_indices_tensor_p,
1086 - num_prefills=num_prefills,
1087 - num_decode_tokens=num_decode_tokens,
1088 - num_prefill_tokens=num_prefill_tokens,
1089 - )
1090 - )
1091 - conv_out_non_spec = torch.vstack(non_spec_parts)
1092 - else:
1093 - conv_out_non_spec = self._short_conv_dilated_decode_batched(
1094 - x_d=x_non_spec,
721 +
722 + if has_decode:
723 + ple_conv(
724 + inputs=inputs_d,
725 + residual=residual_d,
726 conv_state=conv_state,
727 conv_weights=conv_weights,
1097 - state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)],
1098 - has_initial_states_d=metadata.has_initial_states_d,
728 + state_indices=state_indices_d,
729 + mode="decode",
730 + dilation=self.short_conv_dilation,
731 + has_initial_states=metadata.has_initial_states_d,
732 + token_indices=token_indices_d,
733 )
734
1101 - # 3. Merge both parts back into the original token order.
1102 - if has_spec and conv_out_non_spec is not None:
1103 - assert metadata.spec_token_indx is not None
1104 - assert metadata.non_spec_token_indx is not None
1105 - assert spec_output is not None
1106 - output = x.new_empty((metadata.num_actual_tokens, x.size(-1)))
1107 - output.index_copy_(0, metadata.spec_token_indx, spec_output)
1108 - output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec)
1109 - return output
1110 - elif has_spec:
1111 - assert spec_output is not None
1112 - return spec_output
1113 - if conv_out_non_spec is None:
1114 - return x
1115 - return conv_out_non_spec
1116 -
1117 - def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor:
735 + query_start_loc = metadata.non_spec_query_start_loc
736 + if query_start_loc is None:
737 + raise ValueError("query_start_loc is required for prefill short-conv")
738 + query_start_loc = query_start_loc[-num_prefills - 1 :] - num_decode_tokens
739 + has_initial_states = metadata.has_initial_states_p
740 + if has_initial_states is None:
741 + raise ValueError(
742 + "has_initial_states_p is required for prefill short-conv"
743 + )
744 + ple_conv(
745 + inputs=inputs_p,
746 + residual=residual_p,
747 + conv_state=conv_state,
748 + conv_weights=conv_weights,
749 + state_indices=state_indices_p,
750 + mode="prefill",
751 + dilation=self.short_conv_dilation,
752 + query_start_loc=query_start_loc,
753 + has_initial_states=has_initial_states,
754 + token_indices=token_indices_p,
755 + )
756 + else:
757 + num_decode_rows = (
758 + non_spec_token_indices.numel()
759 + if non_spec_token_indices is not None
760 + else inputs.size(0)
761 + )
762 + ple_conv(
763 + inputs=inputs,
764 + residual=residual,
765 + conv_state=conv_state,
766 + conv_weights=conv_weights,
767 + state_indices=state_indices[:num_decode_rows],
768 + mode="decode",
769 + dilation=self.short_conv_dilation,
770 + has_initial_states=metadata.has_initial_states_d,
771 + token_indices=non_spec_token_indices,
772 + )
773 +
774 + def _short_conv(self, inputs: torch.Tensor, residual: torch.Tensor) -> None:
775 forward_context = get_forward_context()
776 attn_metadata = forward_context.attn_metadata
777 + # Profiling omits all metadata or this Mamba entry. The residual
778 + # already contains the gated output, so short convolution is a no-op.
779 if attn_metadata is None:
1121 - return self._short_conv_fallback(inputs)
780 + return
781
782 if not isinstance(attn_metadata, dict):
783 raise RuntimeError(
@@ -1128,8 +787,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
787
788 layer_attn_metadata = attn_metadata.get(self.prefix)
789 if layer_attn_metadata is None:
1131 - # MRV2 omits Mamba-family metadata during profile warmup.
1132 - return self._short_conv_fallback(inputs)
790 + return
791 if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
792 raise TypeError(
793 "Expected PleShortConvAttentionMetadata for layer "
@@ -1138,24 +796,27 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
796 )
797
798 conv_state = self.kv_cache[0]
799 + # Canonicalize both backend cache layouts to [slot, channel, window].
800 if not is_conv_state_dim_first():
801 conv_state = conv_state.transpose(-1, -2)
802 conv_weights = self.conv1d.weight.squeeze(1)
803
804 state_capacity = self.conv_state_len + self.num_spec_tokens
805 if state_capacity > 0:
1147 - if conv_state.size(-1) < state_capacity:
806 + state_size = conv_state.size(-1)
807 + if state_size < state_capacity:
808 raise RuntimeError(
809 "PLE short-conv cache is smaller than expected for "
1150 - f"dilated convolution: got {conv_state.size(-1)}, "
810 + f"dilated convolution: got {state_size}, "
811 f"expect at least {state_capacity}."
812 )
813 conv_state = conv_state[..., -state_capacity:]
1154 - return self._short_conv_dilated_dispatch(
1155 - inputs,
1156 - layer_attn_metadata,
1157 - conv_state,
1158 - conv_weights.to(dtype=inputs.dtype),
814 + self._short_conv_dilated_dispatch(
815 + inputs=inputs,
816 + residual=residual,
817 + metadata=layer_attn_metadata,
818 + conv_state=conv_state,
819 + conv_weights=conv_weights.to(dtype=inputs.dtype),
820 )
821
822 def forward(
@@ -1174,24 +835,21 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
835 )
836 embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
837 embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype)
1177 - key, _ = self.key_proj(embeddings)
1178 - value, _ = self.value_proj(embeddings)
1179 - token_count = hidden_states.shape[0]
1180 - key = key.reshape(token_count, self.hc_count, self.hidden_size)
1181 - query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size)
1182 - key = self._apply_norm(self.norm_key, key)
1183 - query = self._apply_norm(self.norm_query, query)
1184 - gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size)
1185 - gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt())
1186 - gated_value = gate * value.unsqueeze(-2)
1187 - normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2)
1188 - conv_output = torch.zeros_like(normalized)
1189 - torch.ops.vllm.qwen4_exp_ple_short_conv(
1190 - normalized,
1191 - conv_output,
1192 - self.prefix,
1193 - )
1194 - return gated_value.flatten(-2) + conv_output
838 + kv, _ = self.kv_proj(embeddings)
839 + key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
840 + gated_output, conv_input = ple_gate(
841 + key,
842 + value,
843 + hidden_states,
844 + self.norm_key.weight,
845 + self.norm_query.weight,
846 + self.norm_conv.weight,
847 + self.norm_key.eps,
848 + )
849 + # State routing depends on runtime request metadata and remains outside
850 + # the piecewise graph; short convolution accumulates into gated_output.
851 + torch.ops.vllm.qwen4_exp_ple_short_conv(conv_input, gated_output, self.prefix)
852 + return gated_output
853
854
855 def qwen4_exp_compute_ple_ngram_ids(
@@ -1203,12 +861,12 @@ def qwen4_exp_compute_ple_ngram_ids(
861 ) -> None:
862 """Compute request-dependent PLE n-gram IDs outside piecewise graphs."""
863 layer = get_forward_context().no_compile_layers[layer_name]
1206 - ngram_ids = layer.ple_embedding.compute_ngram_ids(
864 + layer.ple_embedding.compute_ngram_ids(
865 input_ids,
866 query_start_loc,
867 ngram_context,
868 + output,
869 )
1211 - output.copy_(ngram_ids)
870
871
872 def qwen4_exp_compute_ple_ngram_ids_fake(
@@ -1223,17 +881,16 @@ def qwen4_exp_compute_ple_ngram_ids_fake(
881
882 def qwen4_exp_ple_short_conv(
883 inputs: torch.Tensor,
1226 - output: torch.Tensor,
884 + residual_output: torch.Tensor,
885 layer_name: str,
886 ) -> None:
887 layer = get_forward_context().no_compile_layers[layer_name]
1230 - result = layer._short_conv(inputs)
1231 - output[: result.shape[0]].copy_(result)
888 + layer._short_conv(inputs, residual_output)
889
890
891 def qwen4_exp_ple_short_conv_fake(
892 inputs: torch.Tensor,
1236 - output: torch.Tensor,
893 + residual_output: torch.Tensor,
894 layer_name: str,
895 ) -> None:
896 return
@@ -1250,7 +907,7 @@ direct_register_custom_op(
907 direct_register_custom_op(
908 op_name="qwen4_exp_ple_short_conv",
909 op_func=qwen4_exp_ple_short_conv,
1253 - mutates_args=["output"],
910 + mutates_args=["residual_output"],
911 fake_impl=qwen4_exp_ple_short_conv_fake,
912 )
913
vllm/v1/attention/backends/short_conv_attn.py
-6
@@ -141,12 +141,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
141 self.spec_sequence_masks = torch.empty(
142 (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device
143 )
144 - self.spec_token_indx = torch.empty(
145 - (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device
146 - )
147 - self.non_spec_token_indx = torch.empty(
148 - (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device
149 - )
144 self.spec_query_start_loc = torch.empty(
145 (self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device
146 )