class Qwen4ExpPLELayer(nn.Module, MambaBase):
def __init__(
self,
config: Qwen4ExpTextConfig,
vllm_config: VllmConfig,
layer_idx: int = 0,
ple_dense_layer_id: int | None = None,
prefix: str = "",
use_sequence_parallel: bool = False,
) -> None:
super().__init__()
self.use_sequence_parallel = use_sequence_parallel
model_config = vllm_config.model_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
self.model_config: ModelConfig = model_config
self.cache_config: CacheConfig = cache_config
self.layer_idx = layer_idx
self.ple_dense_layer_id = (
int(ple_dense_layer_id)
if ple_dense_layer_id is not None
else int(layer_idx)
)
self.prefix = prefix
self.hidden_size = int(config.hidden_size)
self.hc_count = config.hc_count
self.hc_hidden_size = self.hidden_size * self.hc_count
self.conv_kernel_size = int(config.ple_conv_kernel_size)
self.short_conv_dilation = int(config.ngram_size)
self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation
# The conv window ending at a checkpoint is a copy of its last inputs.
self._checkpoint_exporter = MambaPrefillCheckpointExporter(
state_len=self.conv_state_len
)
self.num_spec_tokens = vllm_config.num_speculative_tokens
self.activation = "silu"
self.ple_embedding = Qwen4ExpNGramEmbedding(
config,
int(config.ple_embed_dim),
self.ple_dense_layer_id,
vllm_config.scheduler_config.max_num_batched_tokens,
data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
prefix=f"{prefix}.ple_embedding",
quant_config=quant_config,
params_dtype=model_config.dtype,
use_sequence_parallel=self.use_sequence_parallel,
)
# Projection weights stay replicated even when their token rows are sharded.
self.kv_proj = MergedColumnParallelLinear(
int(config.ple_embed_dim),
[self.hc_hidden_size, self.hidden_size],
bias=False,
params_dtype=model_config.dtype,
quant_config=quant_config,
prefix=f"{prefix}.kv_proj",
disable_tp=True,
)
norm_args = (
self.hc_hidden_size,
config.rms_norm_eps,
self.hidden_size,
model_config.dtype,
)
self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args)
self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args)
self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args)
self.conv1d = nn.Conv1d(
self.hc_hidden_size,
self.hc_hidden_size,
self.conv_kernel_size,
groups=self.hc_hidden_size,
padding=self.conv_state_len,
dilation=self.short_conv_dilation,
bias=False,
dtype=model_config.dtype,
)
nn.init.zeros_(self.conv1d.weight)
self.conv1d.weight._no_reinit = True
self.kv_cache = (torch.tensor([]),)
compilation_config = get_current_vllm_config().compilation_config
if prefix in compilation_config.static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
def _dequantize_embeddings(
self,
embeddings: torch.Tensor,
output_dtype: torch.dtype,
) -> torch.Tensor:
"""Dequantize PLE lookup output."""
return self.ple_embedding.ngram_embedding.dequantize(
embeddings,
output_dtype,
)
def start_prefetch(
self,
input_ids: torch.Tensor,
query_start_loc: torch.Tensor,
ngram_context: torch.Tensor,
) -> None:
"""Start the pinned PLE lookup while the preceding decoder layer runs."""
self.ple_embedding.start_prefetch(input_ids, query_start_loc, ngram_context)
@property
def mamba_type(self) -> MambaAttentionBackendEnum:
return MambaAttentionBackendEnum.SHORT_CONV
@property
def is_kv_cache_tp_replicated(self) -> bool:
return True
def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
return PleShortConvAttentionBackend
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
return MambaStateDtypeCalculator.short_conv_state_dtype(
self.model_config.dtype, self.cache_config.mamba_cache_dtype
)
def get_state_shape(self) -> Sequence[tuple[int, ...]]:
return MambaStateShapeCalculator.short_conv_state_shape(
tp_world_size=1,
intermediate_size=self.hc_hidden_size,
conv_kernel=self.conv_state_len + 1,
num_spec=self.num_spec_tokens,
)
def get_prefill_checkpoint_alignment(self) -> int | None:
return 1 if self.cache_config.mamba_cache_mode == "align" else None
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
spec = super().get_kv_cache_spec(vllm_config)
# The scheduler takes internal checkpoints only when every Mamba layer
# can export them.
if (
self.get_prefill_checkpoint_alignment() is None
or get_mamba_prefill_checkpoint_alignment(vllm_config) is None
or not isinstance(spec, MambaSpec)
):
return spec
grid = get_mamba_prefill_checkpoint_grid(vllm_config)
# Offsets are multiples of the grid, so the window holds only inputs.
assert grid >= self.conv_state_len
return replace(
spec, num_prefill_checkpoint_blocks=1, prefill_checkpoint_alignment=grid
)
def _short_conv_dilated_dispatch(
self,
inputs: torch.Tensor,
residual: torch.Tensor,
outer_residual: torch.Tensor | None,
metadata: PleShortConvAttentionMetadata,
conv_state: torch.Tensor,
conv_weights: torch.Tensor,
) -> None:
num_prefills = metadata.num_prefills
num_decodes = metadata.num_decodes
num_decode_tokens = metadata.num_decode_tokens
num_prefill_tokens = metadata.num_prefill_tokens
has_prefill = num_prefills > 0
has_decode = num_decodes > 0
has_spec = metadata.spec_sequence_masks is not None
has_non_spec = has_prefill or has_decode
inputs = inputs[: metadata.num_actual_tokens]
residual = residual[: metadata.num_actual_tokens]
if outer_residual is not None:
outer_residual = outer_residual[: metadata.num_actual_tokens]
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
if has_spec:
assert metadata.spec_state_indices_tensor is not None
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,
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,
)
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
)
residual_d, residual_p = torch.split(
residual, [num_decode_tokens, num_prefill_tokens], dim=0
)
outer_residual_d: torch.Tensor | None = None
outer_residual_p: torch.Tensor | None = None
if outer_residual is not None:
outer_residual_d, outer_residual_p = torch.split(
outer_residual,
[num_decode_tokens, num_prefill_tokens],
dim=0,
)
token_indices_d = None
token_indices_p = None
else:
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=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,
)
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,
)
if metadata.checkpoint is not None:
assert metadata.non_spec_query_start_loc is not None
self._checkpoint_exporter.export(
metadata.checkpoint,
raw_qkv=inputs,
conv_state=conv_state,
cu_seqlens=metadata.non_spec_query_start_loc,
token_indices=non_spec_token_indices,
)
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,
)
# State routing consumes the current request metadata on every replay.
@eager_break_during_capture
def _short_conv(
self,
inputs: torch.Tensor,
residual: torch.Tensor,
outer_residual: torch.Tensor | None,
) -> None:
"""Update convolution state, optionally fusing both residual additions."""
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:
if outer_residual is not None:
residual.add_(outer_residual)
return
if not isinstance(attn_metadata, dict):
raise RuntimeError(
"PLE short-conv expects per-layer attention metadata dict "
f"during inference, got {type(attn_metadata).__name__}."
)
layer_attn_metadata = attn_metadata.get(self.prefix)
if layer_attn_metadata is None:
if outer_residual is not None:
residual.add_(outer_residual)
return
if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
raise TypeError(
"Expected PleShortConvAttentionMetadata for layer "
f"'{self.prefix}', got "
f"{type(layer_attn_metadata).__name__}."
)
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)
state_capacity = self.conv_state_len + self.num_spec_tokens
if state_capacity > 0:
state_size = conv_state.size(-1)
if state_size < state_capacity:
raise RuntimeError(
"PLE short-conv cache is smaller than expected for "
f"dilated convolution: got {state_size}, "
f"expect at least {state_capacity}."
)
conv_state = conv_state[..., -state_capacity:]
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(
self,
hidden_states: torch.Tensor,
input_ids: torch.Tensor,
query_start_loc: torch.Tensor,
ngram_context: torch.Tensor,
) -> torch.Tensor:
input_ids = input_ids.reshape(-1)
if (
not self.use_sequence_parallel
and input_ids.shape[0] != hidden_states.shape[0]
):
raise ValueError(
"PLE expects input_ids and hidden_states to have the same "
f"token length, got {input_ids.shape[0]} and "
f"{hidden_states.shape[0]}"
)
embeddings = self.ple_embedding(
hidden_states,
input_ids,
query_start_loc,
ngram_context,
)
assert embeddings.shape[0] == hidden_states.shape[0]
embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype)
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,
)
if self.use_sequence_parallel:
# Only convolution needs full tokens. Add residuals on local shards.
conv_input = sp_all_gather(conv_input)[: input_ids.shape[0]]
# Metadata may omit padding tokens or all tokens during profiling.
conv_output = torch.zeros_like(conv_input)
self._short_conv(conv_input, conv_output, outer_residual=None)
gated_output.add_(sp_shard(conv_output))
gated_output.add_(hidden_states)
else:
self._short_conv(conv_input, gated_output, hidden_states)
return gated_output