[ROCm][Qwen4Exp] Use fused PLE Triton kernels on the AMD backend (#60021)
Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Mikko Tukiainen
co-authored by
Cursor
Cursor
parent
0fc0cefb4e
commit
e2414ec1db
@@ -45,6 +45,7 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
|
||||
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
||||
)
|
||||
from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import weak_ref_tensor
|
||||
from vllm.v1.attention.backends.short_conv_attn import (
|
||||
PleShortConvAttentionMetadata,
|
||||
@@ -1102,7 +1103,7 @@ def test_fused_ngram_ids_correctness(
|
||||
contexts: list[list[int]],
|
||||
first_token_id: int,
|
||||
) -> None:
|
||||
from vllm.models.qwen4_exp.nvidia.ops.ple import ple_ngram_ids
|
||||
from vllm.models.qwen4_exp.common.ops.ple import ple_ngram_ids
|
||||
|
||||
device = torch.device("cuda")
|
||||
query_start_loc = torch.tensor(
|
||||
@@ -1857,15 +1858,20 @@ def _make_conv_metadata(
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"layer_cls", [Qwen4ExpPLELayer, Qwen4ExpPLELayerAMD], ids=["nvidia", "amd"]
|
||||
)
|
||||
def test_fused_conv_correctness(
|
||||
case: _ConvBatchCase,
|
||||
state_layout: str,
|
||||
layer_cls: type[nn.Module],
|
||||
) -> None:
|
||||
device = torch.device("cuda")
|
||||
metadata, num_real_tokens = _make_conv_metadata(case, device)
|
||||
if case.state_index_stride > 1:
|
||||
assert metadata.state_indices_tensor.stride(0) == case.state_index_stride
|
||||
module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
|
||||
# layer_cls is type[nn.Module], so layer_cls.__new__ is type.__new__.
|
||||
module = object.__new__(layer_cls)
|
||||
nn.Module.__init__(module)
|
||||
module.conv_state_len = (case.kernel_size - 1) * case.dilation
|
||||
module.short_conv_dilation = case.dilation
|
||||
@@ -1938,7 +1944,7 @@ def test_fused_conv_correctness(
|
||||
def test_fused_gate_correctness(num_tokens: int, strided_kv: bool) -> None:
|
||||
import math
|
||||
|
||||
from vllm.models.qwen4_exp.nvidia.ops.ple import ple_gate
|
||||
from vllm.models.qwen4_exp.common.ops.ple import ple_gate
|
||||
|
||||
device = torch.device("cuda")
|
||||
hc, h = 4, 2560
|
||||
@@ -2038,7 +2044,6 @@ def _build_amd_ngram_embedding(
|
||||
embedding_dim=config.ple_embed_dim,
|
||||
ple_dense_layer_id=0,
|
||||
max_total_tokens=4,
|
||||
max_num_reqs=2,
|
||||
prefix="test.ple_embedding",
|
||||
layer_name="test.ple",
|
||||
quant_config=(
|
||||
@@ -2224,3 +2229,47 @@ def test_amd_pinned_embedding_output_written_under_cudagraph_capture(
|
||||
|
||||
expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2)
|
||||
torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="only the ROCm backend runs the n-gram embedding under torch.compile",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("query_lens", "contexts"),
|
||||
[([3, 1], [[5, 6], [0, 7]]), ([1, 0, 2], [[1, 2], [3, 4], [0, 0]])],
|
||||
ids=["two-requests", "empty-request"],
|
||||
)
|
||||
def test_amd_ngram_embedding_matches_reference_under_compile(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
query_lens: list[int],
|
||||
contexts: list[list[int]],
|
||||
) -> None:
|
||||
"""The AMD n-gram ID op must run under Dynamo with a symbolic request count."""
|
||||
module, loaded_weight = _build_amd_ngram_embedding(
|
||||
monkeypatch, device="cuda:0", cpu_offload=False, fp8_checkpoint=False
|
||||
)
|
||||
device = torch.device("cuda:0")
|
||||
query_start_loc = torch.tensor(
|
||||
[0, *accumulate(query_lens)], dtype=torch.int32, device=device
|
||||
)
|
||||
input_ids = torch.tensor(
|
||||
[9, 0, 11, 12][: sum(query_lens)], dtype=torch.int32, device=device
|
||||
)
|
||||
ngram_context = torch.tensor(contexts, dtype=torch.int32, device=device)
|
||||
|
||||
compiled = torch.compile(module, fullgraph=True, dynamic=True)
|
||||
output = compiled(input_ids, query_start_loc, ngram_context)
|
||||
|
||||
ngram_ids = _reference_ngram_ids(
|
||||
input_ids,
|
||||
query_start_loc,
|
||||
ngram_context,
|
||||
module.layer_multipliers,
|
||||
module.ngram_heads_vocab_sizes,
|
||||
module.ngram_heads_offsets,
|
||||
module.eos_token_id,
|
||||
module.heads_per_ngram,
|
||||
)
|
||||
expected = loaded_weight[ngram_ids.cpu()].flatten(-2).to(device)
|
||||
torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0)
|
||||
|
||||
@@ -140,9 +140,9 @@ _QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [
|
||||
"_input_scale",
|
||||
]
|
||||
|
||||
# The checkpoint keeps down and injection projections separate; runtime packs
|
||||
# them into adjacent logical shards of one MergedColumnParallelLinear.
|
||||
_HC_WEIGHTS_MAPPER = WeightsMapper(
|
||||
# The checkpoint stores these projections separately; runtime packs each group
|
||||
# into adjacent logical shards of a MergedColumnParallelLinear.
|
||||
_EXTRA_WEIGHTS_MAPPER = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
"hyper_connection.input_mix_weight_down.weight": (
|
||||
"hyper_connection.input_mix_weight_down_block_inject.weight",
|
||||
@@ -152,6 +152,8 @@ _HC_WEIGHTS_MAPPER = WeightsMapper(
|
||||
"hyper_connection.input_mix_weight_down_block_inject.weight",
|
||||
1,
|
||||
),
|
||||
"ple.key_proj": ("ple.kv_proj", 0),
|
||||
"ple.value_proj": ("ple.kv_proj", 1),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -276,7 +278,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
|
||||
|
||||
if input_ids is None or query_start_loc is None or ngram_context is None:
|
||||
raise RuntimeError("PLE inputs were not prepared")
|
||||
hidden_states = hidden_states + self.ple(
|
||||
hidden_states = self.ple(
|
||||
hidden_states,
|
||||
input_ids,
|
||||
query_start_loc,
|
||||
@@ -371,7 +373,7 @@ class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
|
||||
}
|
||||
)
|
||||
class Qwen4ExpModel(nn.Module):
|
||||
hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
|
||||
hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
@@ -594,6 +596,7 @@ class Qwen4ExpForCausalLM(
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
|
||||
"in_proj_ba": ["in_proj_b", "in_proj_a"],
|
||||
"input_mix_weight_down_block_inject": [
|
||||
@@ -644,8 +647,6 @@ class Qwen4ExpForCausalLM(
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: object,
|
||||
) -> torch.Tensor | IntermediateTensors:
|
||||
# Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs
|
||||
# path (query_start_loc / ngram_context) reaches Qwen4ExpModel.
|
||||
return self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
@@ -825,11 +826,12 @@ class Qwen4ExpForConditionalGeneration(
|
||||
requires_raw_input_tokens = True
|
||||
|
||||
packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
|
||||
"kv_proj": ["key_proj", "value_proj"],
|
||||
"input_mix_weight_down_block_inject": [
|
||||
"input_mix_weight_down",
|
||||
"block_inject_weight",
|
||||
"_input_mix_padding",
|
||||
]
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -50,7 +50,7 @@ from vllm.sequence import IntermediateTensors
|
||||
from .hyperconnection import GatedResidual, HyperConnectionConfig
|
||||
from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
|
||||
from .model import (
|
||||
_HC_WEIGHTS_MAPPER,
|
||||
_EXTRA_WEIGHTS_MAPPER,
|
||||
_QWEN4_EXP_IGNORED_MISSING_SUFFIXES,
|
||||
Qwen4ExpDecoderLayer,
|
||||
Qwen4ExpMixtureOfExperts,
|
||||
@@ -172,7 +172,7 @@ def _make_draft_vllm_config(
|
||||
}
|
||||
)
|
||||
class Qwen4ExpMultiTokenPredictor(nn.Module):
|
||||
hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
|
||||
hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
|
||||
super().__init__()
|
||||
|
||||
@@ -2,18 +2,16 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GPU-resident Qwen4Exp position-learning enhancement layers."""
|
||||
|
||||
import math
|
||||
from collections.abc import Iterable, Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from transformers import Qwen4ExpTextConfig
|
||||
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.layers.linear import MergedColumnParallelLinear
|
||||
from vllm.model_executor.layers.mamba.abstract import MambaBase
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import (
|
||||
MambaStateDtypeCalculator,
|
||||
@@ -30,13 +28,13 @@ from vllm.v1.attention.backends.short_conv_attn import (
|
||||
PleShortConvAttentionBackend,
|
||||
PleShortConvAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
from ..common.ngram_embedding import (
|
||||
Qwen4ExpPLEDeviceEmbedding,
|
||||
Qwen4ExpPLEEmbeddingMethod,
|
||||
Qwen4ExpPLEPinnedHostEmbedding,
|
||||
)
|
||||
from ..common.ops.ple import ple_conv, ple_gate, ple_ngram_ids
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -177,7 +175,6 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
embedding_dim: int,
|
||||
ple_dense_layer_id: int,
|
||||
max_total_tokens: int,
|
||||
max_num_reqs: int,
|
||||
prefix: str,
|
||||
layer_name: str,
|
||||
*,
|
||||
@@ -270,55 +267,6 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
self.ngram_embedding.weight.device,
|
||||
self.ngram_embedding.weight.is_pinned(),
|
||||
)
|
||||
self.register_buffer(
|
||||
"positions_buffer",
|
||||
torch.arange(max_total_tokens, dtype=torch.int64),
|
||||
persistent=False,
|
||||
)
|
||||
self.register_buffer(
|
||||
"padded_buffer",
|
||||
torch.full(
|
||||
(max_num_reqs, max_total_tokens),
|
||||
self.eos_token_id,
|
||||
dtype=torch.int64,
|
||||
),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _shift_precompute(
|
||||
tokens: torch.Tensor, eos_token_id: int
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if tokens.dim() != 2:
|
||||
raise ValueError("tokens must be a 2D tensor")
|
||||
batch_size, seq_len = tokens.shape
|
||||
positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
|
||||
eos_positions = torch.where(tokens == eos_token_id, positions, -1)
|
||||
previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
|
||||
previous_eos = torch.cat(
|
||||
[
|
||||
eos_positions.new_full((batch_size, 1), -1),
|
||||
previous_eos_inclusive[:, :-1],
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
return positions, positions.unsqueeze(0) - previous_eos - 1
|
||||
|
||||
@staticmethod
|
||||
def _shift_apply(
|
||||
tokens: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
position_in_segment: torch.Tensor,
|
||||
shift: int,
|
||||
eos_token_id: int,
|
||||
) -> torch.Tensor:
|
||||
if shift == 0:
|
||||
return tokens
|
||||
source = positions - shift
|
||||
gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
|
||||
shifted = tokens.gather(1, gather_indices)
|
||||
valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
|
||||
return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -326,64 +274,17 @@ class Qwen4ExpNGramEmbedding(nn.Module):
|
||||
query_start_loc: torch.Tensor,
|
||||
ngram_context: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
input_ids = input_ids.reshape(-1).long()
|
||||
query_start_loc = query_start_loc.long()
|
||||
num_reqs = query_start_loc.numel() - 1
|
||||
num_tokens = input_ids.shape[0]
|
||||
if num_tokens > self.positions_buffer.numel():
|
||||
raise ValueError(
|
||||
f"PLE received {num_tokens} tokens, but its workspace supports "
|
||||
f"at most {self.positions_buffer.numel()}"
|
||||
)
|
||||
if num_reqs > self.padded_buffer.shape[0]:
|
||||
raise ValueError(
|
||||
f"PLE received {num_reqs} requests, but its workspace supports "
|
||||
f"at most {self.padded_buffer.shape[0]}"
|
||||
)
|
||||
|
||||
positions = self.positions_buffer[:num_tokens]
|
||||
packed = self.padded_buffer[:num_reqs]
|
||||
packed.fill_(self.eos_token_id)
|
||||
request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
|
||||
request_indices.clamp_(max=num_reqs - 1)
|
||||
columns = (positions - query_start_loc[request_indices]).clamp(
|
||||
0, packed.shape[1] - 1
|
||||
input_ids = input_ids.reshape(-1)
|
||||
ngram_ids = input_ids.new_empty(
|
||||
(input_ids.shape[0], self.ngram_heads), dtype=torch.int64
|
||||
)
|
||||
packed[request_indices, columns] = input_ids
|
||||
ngram_context = ngram_context[:num_reqs].to(
|
||||
device=input_ids.device, dtype=torch.long
|
||||
torch.ops.vllm.qwen4_exp_amd_ple_ngram_ids(
|
||||
input_ids,
|
||||
query_start_loc,
|
||||
ngram_context,
|
||||
ngram_ids,
|
||||
self.layer_name,
|
||||
)
|
||||
|
||||
context = torch.cat([ngram_context, packed], dim=-1)
|
||||
positions_2d, position_in_segment = self._shift_precompute(
|
||||
context, self.eos_token_id
|
||||
)
|
||||
shifted = [context]
|
||||
for shift in range(1, self.ngram_size):
|
||||
shifted.append(
|
||||
self._shift_apply(
|
||||
context,
|
||||
positions_2d,
|
||||
position_in_segment,
|
||||
shift,
|
||||
self.eos_token_id,
|
||||
)
|
||||
)
|
||||
adjusted_columns = columns + self.ngram_size - 1
|
||||
id_blocks = []
|
||||
for ngram in range(2, self.ngram_size + 1):
|
||||
start = (ngram - 2) * self.heads_per_ngram
|
||||
end = start + self.heads_per_ngram
|
||||
mixed = shifted[0] * self.layer_multipliers[0]
|
||||
for index in range(1, ngram):
|
||||
mixed = torch.bitwise_xor(
|
||||
mixed, shifted[index] * self.layer_multipliers[index]
|
||||
)
|
||||
sizes = self.ngram_heads_vocab_sizes[start:end]
|
||||
offsets = self.ngram_heads_offsets[start:end]
|
||||
ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
|
||||
id_blocks.append(ids[request_indices, adjusted_columns])
|
||||
ngram_ids = torch.cat(id_blocks, dim=-1)
|
||||
embedding = self.ngram_embedding
|
||||
if embedding.supports_prefetch:
|
||||
output = ngram_ids.new_empty(
|
||||
@@ -508,26 +409,21 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
int(config.ple_embed_dim),
|
||||
self.ple_dense_layer_id,
|
||||
vllm_config.scheduler_config.max_num_batched_tokens,
|
||||
vllm_config.scheduler_config.max_num_seqs,
|
||||
f"{prefix}.ple_embedding",
|
||||
prefix,
|
||||
data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
|
||||
quant_config=quant_config,
|
||||
params_dtype=model_config.dtype,
|
||||
)
|
||||
self.key_proj = ReplicatedLinear(
|
||||
# The PLE cache is TP-replicated, so this merged projection is too.
|
||||
self.kv_proj = MergedColumnParallelLinear(
|
||||
int(config.ple_embed_dim),
|
||||
self.hc_hidden_size,
|
||||
[self.hc_hidden_size, self.hidden_size],
|
||||
bias=False,
|
||||
params_dtype=model_config.dtype,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.key_proj",
|
||||
)
|
||||
self.value_proj = ReplicatedLinear(
|
||||
int(config.ple_embed_dim),
|
||||
self.hidden_size,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.value_proj",
|
||||
prefix=f"{prefix}.kv_proj",
|
||||
disable_tp=True,
|
||||
)
|
||||
norm_args = (
|
||||
self.hc_hidden_size,
|
||||
@@ -580,355 +476,15 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
num_spec=self.num_spec_tokens,
|
||||
)
|
||||
|
||||
def _apply_norm(
|
||||
self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
shape = hidden_states.shape
|
||||
return norm(hidden_states.flatten(-2)).reshape(shape)
|
||||
|
||||
def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
# Profiling / CUDA graph capture only; conv state is not updated.
|
||||
inputs_t = inputs.transpose(0, 1).unsqueeze(0)
|
||||
output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)]
|
||||
return F.silu(output).squeeze(0).transpose(0, 1)
|
||||
|
||||
def _short_conv_dilated_decode_batched(
|
||||
self,
|
||||
x_d: torch.Tensor,
|
||||
conv_state: torch.Tensor,
|
||||
conv_weights: torch.Tensor,
|
||||
state_indices_tensor_d: torch.Tensor,
|
||||
has_initial_states_d: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
state_indices = state_indices_tensor_d.to(
|
||||
device=conv_state.device, dtype=torch.int64
|
||||
)
|
||||
# FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
|
||||
# slot 0 for a safe gather, then zero output and skip write-back.
|
||||
valid_state = state_indices != NULL_BLOCK_ID
|
||||
state_indices = torch.where(
|
||||
valid_state, state_indices, torch.zeros_like(state_indices)
|
||||
)
|
||||
if has_initial_states_d is None:
|
||||
has_initial_state = valid_state
|
||||
else:
|
||||
if has_initial_states_d.numel() < state_indices_tensor_d.numel():
|
||||
raise ValueError(
|
||||
"has_initial_states_d size mismatch: "
|
||||
f"got {has_initial_states_d.numel()}, "
|
||||
f"need >= {state_indices_tensor_d.numel()}."
|
||||
)
|
||||
has_initial_state = has_initial_states_d[
|
||||
: state_indices_tensor_d.numel()
|
||||
].to(device=conv_state.device, dtype=torch.bool)
|
||||
has_initial_state = has_initial_state & valid_state
|
||||
|
||||
cached_state = conv_state.index_select(0, state_indices)
|
||||
state = cached_state[..., : self.conv_state_len].to(x_d.dtype)
|
||||
if self.conv_state_len > 0:
|
||||
initial_state = torch.where(
|
||||
has_initial_state.view(-1, 1, 1),
|
||||
state,
|
||||
torch.zeros_like(state),
|
||||
)
|
||||
history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
|
||||
else:
|
||||
history = x_d.unsqueeze(-1)
|
||||
|
||||
conv_output = F.conv1d(
|
||||
history,
|
||||
conv_weights.unsqueeze(1).contiguous(),
|
||||
groups=history.size(1),
|
||||
dilation=self.short_conv_dilation,
|
||||
).squeeze(-1)
|
||||
output = F.silu(conv_output)
|
||||
output = output * valid_state.view(-1, 1).to(output.dtype)
|
||||
|
||||
if self.conv_state_len > 0:
|
||||
next_state = history[..., -self.conv_state_len :]
|
||||
# Padded rows are remapped to the reserved null slot. Preserve its
|
||||
# existing value while writing the new states for valid rows.
|
||||
existing_base_state = cached_state[..., : self.conv_state_len]
|
||||
safe_next_state = torch.where(
|
||||
valid_state.view(-1, 1, 1),
|
||||
next_state.to(conv_state.dtype),
|
||||
existing_base_state,
|
||||
)
|
||||
cached_state[..., : self.conv_state_len] = safe_next_state
|
||||
conv_state.index_copy_(0, state_indices, cached_state)
|
||||
|
||||
return output
|
||||
|
||||
def _short_conv_dilated_prefill_batched(
|
||||
self,
|
||||
x_p: torch.Tensor,
|
||||
metadata: PleShortConvAttentionMetadata,
|
||||
conv_state: torch.Tensor,
|
||||
conv_weights: torch.Tensor,
|
||||
state_indices_tensor_p: torch.Tensor,
|
||||
num_prefills: int,
|
||||
num_decode_tokens: int,
|
||||
num_prefill_tokens: int,
|
||||
) -> torch.Tensor:
|
||||
query_start_loc_p = metadata.query_start_loc_p
|
||||
if query_start_loc_p is None:
|
||||
raise ValueError("query_start_loc is required for prefill short-conv")
|
||||
# The metadata builder guarantees that the prefill query offsets start
|
||||
# at 0 and end at num_prefill_tokens. Avoid reading those values here,
|
||||
# since doing so would force a device-to-host synchronization.
|
||||
has_initial_states_p = metadata.has_initial_states_p
|
||||
if has_initial_states_p is None:
|
||||
raise ValueError("has_initial_states_p is required for prefill short-conv")
|
||||
|
||||
output = torch.empty_like(x_p)
|
||||
q_starts = query_start_loc_p.to(torch.int64)
|
||||
if state_indices_tensor_p.numel() < num_prefills:
|
||||
raise ValueError(
|
||||
"state_indices_tensor_p size mismatch: "
|
||||
f"got {state_indices_tensor_p.numel()}, "
|
||||
f"need >= {num_prefills}."
|
||||
)
|
||||
if has_initial_states_p.numel() < num_prefills:
|
||||
raise ValueError(
|
||||
"has_initial_states_p size mismatch: "
|
||||
f"got {has_initial_states_p.numel()}, "
|
||||
f"need >= {num_prefills}."
|
||||
)
|
||||
if num_prefills == 0 or x_p.numel() == 0:
|
||||
return output
|
||||
lengths = q_starts[1:] - q_starts[:-1]
|
||||
# Use the CPU-computed packing width from the metadata builder instead
|
||||
# of synchronizing on lengths.max().
|
||||
max_len = metadata.max_prefill_query_len
|
||||
if max_len <= 0:
|
||||
return output
|
||||
|
||||
hidden_size = x_p.shape[1]
|
||||
positions = torch.arange(
|
||||
num_prefill_tokens, device=x_p.device, dtype=torch.int64
|
||||
)
|
||||
req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
|
||||
col_indices = positions - q_starts[req_indices]
|
||||
|
||||
packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
|
||||
packed_tokens[req_indices, col_indices] = x_p
|
||||
packed_tokens = packed_tokens.transpose(1, 2).contiguous()
|
||||
|
||||
state_indices = state_indices_tensor_p[:num_prefills].to(
|
||||
device=conv_state.device, dtype=torch.int64
|
||||
)
|
||||
valid_state = state_indices != NULL_BLOCK_ID
|
||||
state_indices = torch.where(
|
||||
valid_state, state_indices, torch.zeros_like(state_indices)
|
||||
)
|
||||
has_initial = has_initial_states_p[:num_prefills].to(
|
||||
device=conv_state.device, dtype=torch.bool
|
||||
)
|
||||
if self.conv_state_len > 0:
|
||||
if conv_state.shape[0] == 0:
|
||||
state = conv_state.new_zeros(
|
||||
(num_prefills, hidden_size, self.conv_state_len),
|
||||
dtype=x_p.dtype,
|
||||
)
|
||||
else:
|
||||
state = conv_state.index_select(0, state_indices)[
|
||||
..., : self.conv_state_len
|
||||
].to(x_p.dtype)
|
||||
use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
|
||||
initial_state = torch.where(
|
||||
use_initial_mask,
|
||||
state,
|
||||
torch.zeros_like(state),
|
||||
)
|
||||
history = torch.cat((initial_state, packed_tokens), dim=-1)
|
||||
else:
|
||||
history = packed_tokens
|
||||
|
||||
conv_output = F.conv1d(
|
||||
history,
|
||||
conv_weights.unsqueeze(1).contiguous(),
|
||||
groups=history.size(1),
|
||||
dilation=self.short_conv_dilation,
|
||||
)
|
||||
conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
|
||||
|
||||
token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
|
||||
valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
|
||||
valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
|
||||
num_prefills, 1
|
||||
)
|
||||
conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
|
||||
output.copy_(conv_output[req_indices, col_indices])
|
||||
|
||||
if self.conv_state_len > 0 and conv_state.shape[0] > 0:
|
||||
state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
|
||||
num_prefills, 1, 1
|
||||
)
|
||||
state_offsets = torch.arange(
|
||||
self.conv_state_len, device=history.device, dtype=torch.int64
|
||||
).view(1, 1, self.conv_state_len)
|
||||
next_state = history.gather(
|
||||
dim=2,
|
||||
index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
|
||||
)
|
||||
# Write back without a host synchronization. Valid, non-empty rows
|
||||
# receive their new state; padding and zero-length rows keep the
|
||||
# current cache value.
|
||||
existing_state = conv_state.index_select(0, state_indices)
|
||||
existing_base_state = existing_state[..., : self.conv_state_len]
|
||||
update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
|
||||
safe_next_state = torch.where(
|
||||
update_mask.view(num_prefills, 1, 1),
|
||||
next_state.to(conv_state.dtype),
|
||||
existing_base_state,
|
||||
)
|
||||
existing_state[..., : self.conv_state_len] = safe_next_state
|
||||
conv_state.index_copy_(0, state_indices, existing_state)
|
||||
return output
|
||||
|
||||
def _short_conv_dilated_spec_batched(
|
||||
self,
|
||||
x_spec: torch.Tensor,
|
||||
conv_state: torch.Tensor,
|
||||
conv_weights: torch.Tensor,
|
||||
spec_state_indices_tensor: torch.Tensor,
|
||||
spec_query_start_loc: torch.Tensor,
|
||||
num_accepted_tokens: torch.Tensor,
|
||||
spec_query_len: int,
|
||||
) -> torch.Tensor:
|
||||
"""Dilated short-conv for speculative-decode (MTP) requests.
|
||||
|
||||
Each spec request feeds multiple (draft + 1) query tokens. The conv
|
||||
outputs are computed causally after rolling back the previous draft
|
||||
state by ``num_accepted_tokens - 1``. The current candidate inputs stay
|
||||
in the extended cache for the next forward, matching
|
||||
``causal_conv1d_update``.
|
||||
|
||||
``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
|
||||
length and is a Python int, so no host synchronization is needed; this
|
||||
keeps the path safe for full CUDA-graph capture/replay where the buffers
|
||||
are padded at the request level.
|
||||
"""
|
||||
num_reqs = spec_state_indices_tensor.numel()
|
||||
hidden_size = x_spec.size(-1)
|
||||
# Use a fixed packing width instead of synchronizing on lengths.max().
|
||||
max_len = spec_query_len
|
||||
# Full CUDA graphs can pad these buffers. Only the first num_reqs
|
||||
# accepted-token counts belong to actual speculative requests.
|
||||
num_accepted_tokens = num_accepted_tokens[:num_reqs]
|
||||
q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
|
||||
# Keep the number of real speculative tokens on the device.
|
||||
total_real_tokens = q_starts[num_reqs]
|
||||
|
||||
state_indices = spec_state_indices_tensor.to(
|
||||
device=conv_state.device, dtype=torch.int64
|
||||
)
|
||||
valid_state = state_indices != NULL_BLOCK_ID
|
||||
state_indices = torch.where(
|
||||
valid_state, state_indices, torch.zeros_like(state_indices)
|
||||
)
|
||||
positions = torch.arange(
|
||||
x_spec.size(0), device=x_spec.device, dtype=torch.int64
|
||||
)
|
||||
# Route graph-padded token rows to the discarded dummy request so that
|
||||
# they cannot overwrite real packed data.
|
||||
req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
|
||||
valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
|
||||
clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
|
||||
col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
|
||||
pack_req_indices = torch.where(
|
||||
valid_tokens,
|
||||
clamped_req_indices,
|
||||
torch.full_like(req_indices, num_reqs),
|
||||
)
|
||||
pack_col_indices = torch.where(
|
||||
valid_tokens, col_indices, torch.zeros_like(col_indices)
|
||||
)
|
||||
|
||||
# The last request row is the dummy sink for graph padding.
|
||||
packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
|
||||
packed[pack_req_indices, pack_col_indices] = x_spec
|
||||
packed = packed.transpose(1, 2).contiguous()
|
||||
|
||||
if self.conv_state_len > 0:
|
||||
cached_state = conv_state.index_select(0, state_indices)
|
||||
rollback_offsets = num_accepted_tokens.to(
|
||||
device=conv_state.device, dtype=torch.int64
|
||||
).sub(1)
|
||||
rollback_offsets = torch.where(
|
||||
valid_state,
|
||||
rollback_offsets.clamp_(0, max_len - 1),
|
||||
torch.zeros_like(rollback_offsets),
|
||||
)
|
||||
state_offsets = torch.arange(
|
||||
self.conv_state_len, device=conv_state.device, dtype=torch.int64
|
||||
).view(1, 1, self.conv_state_len)
|
||||
rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
|
||||
state = cached_state.gather(
|
||||
2, rollback_indices.expand(-1, hidden_size, -1)
|
||||
).to(x_spec.dtype)
|
||||
state = torch.where(
|
||||
valid_state.view(num_reqs, 1, 1),
|
||||
state,
|
||||
torch.zeros_like(state),
|
||||
)
|
||||
# Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
|
||||
dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len))
|
||||
state_full = torch.cat((state, dummy_state), dim=0)
|
||||
history = torch.cat((state_full, packed), dim=-1)
|
||||
else:
|
||||
history = packed
|
||||
|
||||
conv_output = F.conv1d(
|
||||
history,
|
||||
conv_weights.unsqueeze(1).contiguous(),
|
||||
groups=history.size(1),
|
||||
dilation=self.short_conv_dilation,
|
||||
)
|
||||
conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
|
||||
|
||||
output = conv_output[pack_req_indices, pack_col_indices]
|
||||
output = output * valid_tokens.view(-1, 1).to(output.dtype)
|
||||
|
||||
# Keep all current candidate inputs in the extended state. On the next
|
||||
# target forward, ``num_accepted_tokens - 1`` selects the rollback
|
||||
# window before processing the newly scheduled tokens.
|
||||
if self.conv_state_len > 0:
|
||||
state_capacity = self.conv_state_len + max_len - 1
|
||||
if conv_state.size(-1) < state_capacity:
|
||||
raise RuntimeError(
|
||||
"PLE short-conv cache cannot retain speculative tokens: "
|
||||
f"got {conv_state.size(-1)}, need {state_capacity}."
|
||||
)
|
||||
candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
|
||||
query_lengths = q_starts[1:] - q_starts[:-1]
|
||||
state_positions = torch.arange(
|
||||
state_capacity, device=history.device, dtype=torch.int64
|
||||
).view(1, 1, state_capacity)
|
||||
update_lengths = (self.conv_state_len + query_lengths - 1).view(
|
||||
num_reqs, 1, 1
|
||||
)
|
||||
update_mask = valid_state.view(num_reqs, 1, 1) & (
|
||||
state_positions < update_lengths
|
||||
)
|
||||
existing_state = cached_state[..., :state_capacity]
|
||||
next_state = torch.where(
|
||||
update_mask,
|
||||
candidate_state.to(conv_state.dtype),
|
||||
existing_state,
|
||||
)
|
||||
cached_state[..., :state_capacity] = next_state
|
||||
conv_state.index_copy_(0, state_indices, cached_state)
|
||||
|
||||
return output
|
||||
|
||||
def _short_conv_dilated_dispatch(
|
||||
self,
|
||||
inputs: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
outer_residual: torch.Tensor,
|
||||
metadata: PleShortConvAttentionMetadata,
|
||||
conv_state: torch.Tensor,
|
||||
conv_weights: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
) -> None:
|
||||
num_prefills = metadata.num_prefills
|
||||
num_decodes = metadata.num_decodes
|
||||
num_decode_tokens = metadata.num_decode_tokens
|
||||
@@ -936,110 +492,145 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
has_prefill = num_prefills > 0
|
||||
has_decode = num_decodes > 0
|
||||
has_spec = metadata.spec_sequence_masks is not None
|
||||
x = inputs[: metadata.num_actual_tokens]
|
||||
has_non_spec = has_prefill or has_decode
|
||||
inputs = inputs[: metadata.num_actual_tokens]
|
||||
residual = residual[: metadata.num_actual_tokens]
|
||||
outer_residual = outer_residual[: metadata.num_actual_tokens]
|
||||
|
||||
# Split spec / non-spec tokens.
|
||||
if has_spec:
|
||||
if has_prefill or has_decode:
|
||||
assert metadata.spec_token_indx is not None
|
||||
assert metadata.non_spec_token_indx is not None
|
||||
x_spec = x.index_select(0, metadata.spec_token_indx.long())
|
||||
x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long())
|
||||
else:
|
||||
x_spec = x
|
||||
x_non_spec = None
|
||||
else:
|
||||
x_spec = None
|
||||
x_non_spec = x
|
||||
spec_token_indices = None
|
||||
non_spec_token_indices = None
|
||||
if has_spec and has_non_spec:
|
||||
assert metadata.spec_token_indx is not None
|
||||
assert metadata.non_spec_token_indx is not None
|
||||
spec_token_indices = metadata.spec_token_indx
|
||||
non_spec_token_indices = metadata.non_spec_token_indx
|
||||
|
||||
spec_output = None
|
||||
# 1. Run the multi-query speculative-decode part.
|
||||
if has_spec:
|
||||
assert metadata.spec_state_indices_tensor is not None
|
||||
assert metadata.spec_query_start_loc is not None
|
||||
assert metadata.num_accepted_tokens is not None
|
||||
spec_output = self._short_conv_dilated_spec_batched(
|
||||
x_spec=x_spec,
|
||||
query_start_loc = metadata.spec_query_start_loc
|
||||
num_accepted_tokens = metadata.num_accepted_tokens
|
||||
assert query_start_loc is not None
|
||||
assert num_accepted_tokens is not None
|
||||
spec_state_indices = metadata.spec_state_indices_tensor[
|
||||
: metadata.num_spec_decodes
|
||||
]
|
||||
# Mixed batches stay in their original row order; the kernels map
|
||||
# logical spec/non-spec rows instead of materializing both groups.
|
||||
ple_conv(
|
||||
inputs=inputs,
|
||||
residual=residual,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
spec_state_indices_tensor=metadata.spec_state_indices_tensor[
|
||||
: metadata.num_spec_decodes
|
||||
],
|
||||
spec_query_start_loc=metadata.spec_query_start_loc,
|
||||
num_accepted_tokens=metadata.num_accepted_tokens,
|
||||
state_indices=spec_state_indices,
|
||||
outer_residual=outer_residual,
|
||||
mode="spec",
|
||||
dilation=self.short_conv_dilation,
|
||||
query_start_loc=query_start_loc,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
spec_query_len=metadata.spec_query_len,
|
||||
token_indices=spec_token_indices,
|
||||
)
|
||||
|
||||
# 2. Run regular decode and prefill requests.
|
||||
conv_out_non_spec = None
|
||||
state_indices_tensor = metadata.state_indices_tensor
|
||||
if x_non_spec is not None:
|
||||
assert state_indices_tensor is not None
|
||||
if has_prefill:
|
||||
state_indices_tensor_d, state_indices_tensor_p = torch.split(
|
||||
state_indices_tensor,
|
||||
[num_decodes, num_prefills],
|
||||
dim=0,
|
||||
if not has_non_spec:
|
||||
return
|
||||
|
||||
state_indices = metadata.state_indices_tensor
|
||||
assert state_indices is not None
|
||||
if has_prefill:
|
||||
state_indices_d, state_indices_p = torch.split(
|
||||
state_indices, [num_decodes, num_prefills], dim=0
|
||||
)
|
||||
if non_spec_token_indices is None:
|
||||
inputs_d, inputs_p = torch.split(
|
||||
inputs, [num_decode_tokens, num_prefill_tokens], dim=0
|
||||
)
|
||||
x_d, x_p = torch.split(
|
||||
x_non_spec,
|
||||
residual_d, residual_p = torch.split(
|
||||
residual, [num_decode_tokens, num_prefill_tokens], dim=0
|
||||
)
|
||||
outer_residual_d, outer_residual_p = torch.split(
|
||||
outer_residual,
|
||||
[num_decode_tokens, num_prefill_tokens],
|
||||
dim=0,
|
||||
)
|
||||
non_spec_parts: list[torch.Tensor] = []
|
||||
if has_decode:
|
||||
non_spec_parts.append(
|
||||
self._short_conv_dilated_decode_batched(
|
||||
x_d=x_d,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
state_indices_tensor_d=state_indices_tensor_d,
|
||||
has_initial_states_d=metadata.has_initial_states_d,
|
||||
)
|
||||
)
|
||||
non_spec_parts.append(
|
||||
self._short_conv_dilated_prefill_batched(
|
||||
x_p=x_p,
|
||||
metadata=metadata,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
state_indices_tensor_p=state_indices_tensor_p,
|
||||
num_prefills=num_prefills,
|
||||
num_decode_tokens=num_decode_tokens,
|
||||
num_prefill_tokens=num_prefill_tokens,
|
||||
)
|
||||
)
|
||||
conv_out_non_spec = torch.vstack(non_spec_parts)
|
||||
token_indices_d = None
|
||||
token_indices_p = None
|
||||
else:
|
||||
conv_out_non_spec = self._short_conv_dilated_decode_batched(
|
||||
x_d=x_non_spec,
|
||||
inputs_d = inputs_p = inputs
|
||||
residual_d = residual_p = residual
|
||||
outer_residual_d = outer_residual_p = outer_residual
|
||||
token_indices_d, token_indices_p = torch.split(
|
||||
non_spec_token_indices,
|
||||
[num_decode_tokens, num_prefill_tokens],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
if has_decode:
|
||||
ple_conv(
|
||||
inputs=inputs_d,
|
||||
residual=residual_d,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)],
|
||||
has_initial_states_d=metadata.has_initial_states_d,
|
||||
state_indices=state_indices_d,
|
||||
outer_residual=outer_residual_d,
|
||||
mode="decode",
|
||||
dilation=self.short_conv_dilation,
|
||||
has_initial_states=metadata.has_initial_states_d,
|
||||
token_indices=token_indices_d,
|
||||
)
|
||||
|
||||
# 3. Merge both parts back into the original token order.
|
||||
if has_spec and conv_out_non_spec is not None:
|
||||
assert metadata.spec_token_indx is not None
|
||||
assert metadata.non_spec_token_indx is not None
|
||||
assert spec_output is not None
|
||||
output = x.new_empty((metadata.num_actual_tokens, x.size(-1)))
|
||||
output.index_copy_(0, metadata.spec_token_indx, spec_output)
|
||||
output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec)
|
||||
return output
|
||||
elif has_spec:
|
||||
assert spec_output is not None
|
||||
return spec_output
|
||||
if conv_out_non_spec is None:
|
||||
return x
|
||||
return conv_out_non_spec
|
||||
query_start_loc = metadata.query_start_loc_p
|
||||
if query_start_loc is None:
|
||||
raise ValueError("query_start_loc is required for prefill short-conv")
|
||||
has_initial_states = metadata.has_initial_states_p
|
||||
if has_initial_states is None:
|
||||
raise ValueError(
|
||||
"has_initial_states_p is required for prefill short-conv"
|
||||
)
|
||||
ple_conv(
|
||||
inputs=inputs_p,
|
||||
residual=residual_p,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
state_indices=state_indices_p,
|
||||
outer_residual=outer_residual_p,
|
||||
mode="prefill",
|
||||
dilation=self.short_conv_dilation,
|
||||
query_start_loc=query_start_loc,
|
||||
has_initial_states=has_initial_states,
|
||||
token_indices=token_indices_p,
|
||||
)
|
||||
else:
|
||||
num_decode_rows = (
|
||||
non_spec_token_indices.numel()
|
||||
if non_spec_token_indices is not None
|
||||
else inputs.size(0)
|
||||
)
|
||||
ple_conv(
|
||||
inputs=inputs,
|
||||
residual=residual,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights,
|
||||
state_indices=state_indices[:num_decode_rows],
|
||||
outer_residual=outer_residual,
|
||||
mode="decode",
|
||||
dilation=self.short_conv_dilation,
|
||||
has_initial_states=metadata.has_initial_states_d,
|
||||
token_indices=non_spec_token_indices,
|
||||
)
|
||||
|
||||
def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor:
|
||||
def _short_conv(
|
||||
self,
|
||||
inputs: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
outer_residual: torch.Tensor,
|
||||
) -> None:
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
# Profiling omits all metadata or this Mamba entry. Short convolution
|
||||
# is a no-op there, but preserve the outer residual addition.
|
||||
if attn_metadata is None:
|
||||
return self._short_conv_fallback(inputs)
|
||||
residual.add_(outer_residual)
|
||||
return
|
||||
|
||||
if not isinstance(attn_metadata, dict):
|
||||
raise RuntimeError(
|
||||
@@ -1049,8 +640,8 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
|
||||
layer_attn_metadata = attn_metadata.get(self.prefix)
|
||||
if layer_attn_metadata is None:
|
||||
# MRV2 omits Mamba-family metadata during profile warmup.
|
||||
return self._short_conv_fallback(inputs)
|
||||
residual.add_(outer_residual)
|
||||
return
|
||||
if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
|
||||
raise TypeError(
|
||||
"Expected PleShortConvAttentionMetadata for layer "
|
||||
@@ -1059,6 +650,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
)
|
||||
|
||||
conv_state = self.kv_cache[0]
|
||||
# Canonicalize both backend cache layouts to [slot, channel, window].
|
||||
if not is_conv_state_dim_first():
|
||||
conv_state = conv_state.transpose(-1, -2)
|
||||
conv_weights = self.conv1d.weight.squeeze(1)
|
||||
@@ -1072,11 +664,13 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
f"expect at least {state_capacity}."
|
||||
)
|
||||
conv_state = conv_state[..., -state_capacity:]
|
||||
return self._short_conv_dilated_dispatch(
|
||||
inputs,
|
||||
layer_attn_metadata,
|
||||
conv_state,
|
||||
conv_weights.to(dtype=inputs.dtype),
|
||||
self._short_conv_dilated_dispatch(
|
||||
inputs=inputs,
|
||||
residual=residual,
|
||||
outer_residual=outer_residual,
|
||||
metadata=layer_attn_metadata,
|
||||
conv_state=conv_state,
|
||||
conv_weights=conv_weights.to(dtype=inputs.dtype),
|
||||
)
|
||||
|
||||
def forward(
|
||||
@@ -1095,27 +689,58 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
)
|
||||
embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
|
||||
embeddings = self.ple_embedding.ngram_embedding.dequantize(
|
||||
embeddings,
|
||||
self.ple_embedding.ngram_embedding.params_dtype,
|
||||
embeddings, hidden_states.dtype
|
||||
)
|
||||
key, _ = self.key_proj(embeddings)
|
||||
value, _ = self.value_proj(embeddings)
|
||||
token_count = hidden_states.shape[0]
|
||||
key = key.reshape(token_count, self.hc_count, self.hidden_size)
|
||||
query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size)
|
||||
key = self._apply_norm(self.norm_key, key)
|
||||
query = self._apply_norm(self.norm_query, query)
|
||||
gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size)
|
||||
gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt())
|
||||
gated_value = gate * value.unsqueeze(-2)
|
||||
normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2)
|
||||
conv_output = torch.zeros_like(normalized)
|
||||
kv, _ = self.kv_proj(embeddings)
|
||||
key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
|
||||
gated_output, conv_input = ple_gate(
|
||||
key,
|
||||
value,
|
||||
hidden_states,
|
||||
self.norm_key.weight,
|
||||
self.norm_query.weight,
|
||||
self.norm_conv.weight,
|
||||
self.norm_key.eps,
|
||||
)
|
||||
# The short-conv op is a piecewise-graph splitting op: state routing
|
||||
# reads runtime request metadata. It accumulates the convolution and
|
||||
# the outer residual into gated_output.
|
||||
torch.ops.vllm.qwen4_exp_ple_short_conv(
|
||||
normalized,
|
||||
conv_output,
|
||||
conv_input,
|
||||
gated_output,
|
||||
hidden_states,
|
||||
self.prefix,
|
||||
)
|
||||
return gated_value.flatten(-2) + conv_output
|
||||
return gated_output
|
||||
|
||||
|
||||
def qwen4_exp_amd_ple_ngram_ids(
|
||||
input_ids: torch.Tensor,
|
||||
query_start_loc: torch.Tensor,
|
||||
ngram_context: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
"""Hash the current request layout into n-gram embedding indices.
|
||||
|
||||
The launch sizes its request binary search from the request count, which
|
||||
is symbolic under Dynamo, so it runs outside Inductor's FX graph.
|
||||
"""
|
||||
layer = get_forward_context().no_compile_layers[layer_name]
|
||||
if not isinstance(layer, Qwen4ExpPLELayer):
|
||||
raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
|
||||
embedding = layer.ple_embedding
|
||||
ple_ngram_ids(
|
||||
input_ids=input_ids,
|
||||
query_start_loc=query_start_loc,
|
||||
ngram_context=ngram_context,
|
||||
layer_multipliers=embedding.layer_multipliers,
|
||||
ngram_heads_vocab_sizes=embedding.ngram_heads_vocab_sizes,
|
||||
ngram_heads_offsets=embedding.ngram_heads_offsets,
|
||||
eos_token_id=embedding.eos_token_id,
|
||||
heads_per_ngram=embedding.heads_per_ngram,
|
||||
output=output,
|
||||
)
|
||||
|
||||
|
||||
def qwen4_exp_amd_ple_ngram_embedding(
|
||||
@@ -1155,12 +780,19 @@ def qwen4_exp_amd_ple_ngram_embedding_pinned(
|
||||
|
||||
def qwen4_exp_ple_short_conv(
|
||||
inputs: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
outer_residual: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> None:
|
||||
layer = get_forward_context().no_compile_layers[layer_name]
|
||||
result = layer._short_conv(inputs)
|
||||
output[: result.shape[0]].copy_(result)
|
||||
layer._short_conv(inputs, residual, outer_residual)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="qwen4_exp_amd_ple_ngram_ids",
|
||||
op_func=qwen4_exp_amd_ple_ngram_ids,
|
||||
mutates_args=["output"],
|
||||
)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
@@ -1180,7 +812,7 @@ direct_register_custom_op(
|
||||
direct_register_custom_op(
|
||||
op_name="qwen4_exp_ple_short_conv",
|
||||
op_func=qwen4_exp_ple_short_conv,
|
||||
mutates_args=["output"],
|
||||
mutates_args=["residual"],
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Qwen4Exp kernels shared by the NVIDIA and AMD backends; import leaf modules
|
||||
directly."""
|
||||
@@ -23,7 +23,7 @@ from ..common.ngram_embedding import (
|
||||
Qwen4ExpPLEPinnedHostEmbedding,
|
||||
Qwen4ExpPLEUnquantizedEmbeddingMethod,
|
||||
)
|
||||
from .ops.ple import ple_ngram_ids
|
||||
from ..common.ops.ple import ple_ngram_ids
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@@ -24,8 +24,8 @@ from vllm.v1.attention.backends.short_conv_attn import (
|
||||
PleShortConvAttentionMetadata,
|
||||
)
|
||||
|
||||
from ..common.ops.ple import ple_conv, ple_gate
|
||||
from .ngram_embedding import Qwen4ExpNGramEmbedding
|
||||
from .ops.ple import ple_conv, ple_gate
|
||||
|
||||
|
||||
class Qwen4ExpPLEGroupedNorm(nn.Module):
|
||||
|
||||
Reference in New Issue
Block a user