class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
metadata_cls = PleShortConvAttentionMetadata
# Spec-decode requires a uniform (multi-token) decode batch for full
# CUDA graph capture, matching the GDN backend.
_cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
reorder_batch_threshold: int = 1
supports_update_block_table = False
needs_causal_conv1d_metadata = False
def __init__(
self,
kv_cache_spec: MambaSpec,
layer_names: list[str],
vllm_config: VllmConfig,
device: torch.device,
) -> None:
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
self.num_spec = self.num_spec_tokens
self.use_full_cuda_graph = (
self.compilation_config.cudagraph_mode.has_full_cudagraphs()
)
max_num_seqs = vllm_config.scheduler_config.max_num_seqs
max_capture_size = self.compilation_config.max_cudagraph_capture_size
self.decode_cudagraph_max_bs = max_num_seqs
self.decode_cudagraph_max_tokens = max_num_seqs * (self.num_spec + 1)
if max_capture_size is not None:
self.decode_cudagraph_max_bs = min(
self.decode_cudagraph_max_bs, max_capture_size
)
self.decode_cudagraph_max_tokens = min(
self.decode_cudagraph_max_tokens, max_capture_size
)
# Persistent buffers reused during full CUDA graph capture and replay.
self.spec_state_indices_tensor = torch.empty(
(self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device
)
self.spec_sequence_masks = torch.empty(
(self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device
)
self.spec_query_start_loc = torch.empty(
(self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device
)
self.num_accepted_tokens = torch.empty(
(self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device
)
self.has_initial_states_d = torch.empty(
(self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device
)
def _build_non_spec_metadata(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool,
num_decode_draft_tokens_cpu: torch.Tensor | None,
**kwargs: Any,
) -> PleShortConvAttentionMetadata:
metadata = super().build(
common_prefix_len,
common_attn_metadata,
fast_build,
num_accepted_tokens=None,
**kwargs,
)
assert isinstance(metadata, PleShortConvAttentionMetadata)
state_indices_d = metadata.state_indices_tensor_d
if state_indices_d is not None and state_indices_d.dim() > 1:
state_indices_d = state_indices_d[:, 0]
state_indices_p = metadata.state_indices_tensor_p
if metadata.num_prefills == 0:
assert state_indices_d is not None
# BaseMambaAttentionMetadataBuilder pads decode state indices into
# a persistent tensor for full CUDA graphs. Keep those rows so the
# PLE decode receives one cache slot per graph-padded token.
state_indices_tensor = state_indices_d
elif metadata.num_decodes == 0:
assert state_indices_p is not None
state_indices_tensor = state_indices_p[: metadata.num_prefills]
else:
assert state_indices_d is not None
assert state_indices_p is not None
state_indices_tensor = torch.cat(
(state_indices_d, state_indices_p[: metadata.num_prefills])
)
has_initial_states_d = None
if metadata.num_decodes > 0:
num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
has_initial_states_d = num_computed_tokens[: metadata.num_decodes] > 0
if (
self.use_full_cuda_graph
and metadata.num_prefills == 0
and metadata.num_decodes <= self.decode_cudagraph_max_bs
):
assert state_indices_d is not None
# Prepare tensors for CUDA graph replay. Padded rows have no
# initial state and use NULL_BLOCK_ID in state_indices_d.
num_decode_rows = state_indices_d.numel()
self.has_initial_states_d[: metadata.num_decodes].copy_(
has_initial_states_d, non_blocking=True
)
self.has_initial_states_d[metadata.num_decodes : num_decode_rows].fill_(
False
)
has_initial_states_d = self.has_initial_states_d[:num_decode_rows]
max_prefill_query_len = 0
if metadata.num_prefills > 0:
query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu)
max_prefill_query_len = int(
query_lens_cpu[
metadata.num_decodes : (
metadata.num_decodes + metadata.num_prefills
)
]
.max()
.item()
)
return replace(
metadata,
num_actual_tokens=common_attn_metadata.num_actual_tokens,
spec_query_len=self.num_spec + 1,
max_prefill_query_len=max_prefill_query_len,
query_start_loc=common_attn_metadata.query_start_loc,
state_indices_tensor=state_indices_tensor,
has_initial_states_d=has_initial_states_d,
non_spec_query_start_loc=common_attn_metadata.query_start_loc,
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
)
def build( # type: ignore[override]
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
*,
num_accepted_tokens: torch.Tensor | None = None,
num_decode_draft_tokens_cpu: torch.Tensor | None = None,
**kwargs: Any,
) -> PleShortConvAttentionMetadata:
m = common_attn_metadata
spec_sequence_masks_cpu: torch.Tensor | None = None
# Detect speculative-decode requests. We use -1 to mark prefill and
# plain-decode requests, so any value >= 0 is a (multi-query)
# spec-decode request.
if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
candidate_mask = num_decode_draft_tokens_cpu[: m.num_reqs] >= 0
if bool(candidate_mask.any().item()):
spec_sequence_masks_cpu = candidate_mask
if spec_sequence_masks_cpu is None:
return self._build_non_spec_metadata(
common_prefix_len,
common_attn_metadata,
fast_build,
num_decode_draft_tokens_cpu,
**kwargs,
)
del common_prefix_len, fast_build, kwargs
query_start_loc = m.query_start_loc
query_start_loc_cpu = m.query_start_loc_cpu
query_lens_cpu = torch.diff(query_start_loc_cpu)
block_table_tensor = mamba_get_block_table_tensor(
m.block_table_tensor,
m.seq_lens,
self.kv_cache_spec,
self.vllm_config.cache_config.mamba_cache_mode,
)
if query_start_loc.device.type == "cpu":
spec_sequence_masks = spec_sequence_masks_cpu
else:
spec_sequence_masks = async_tensor_h2d(
spec_sequence_masks_cpu, device=query_start_loc.device
)
has_initial_states_p = None
has_initial_states_d = None
# Per-request classification by mask, NOT by position. With
# spec-decode, the front decode group can contain both spec-decode
# requests and plain non-spec single-token decodes. Spec requests are
# therefore not guaranteed to occupy the first num_spec_decodes slots.
non_spec_mask_cpu = ~spec_sequence_masks_cpu
decode_mask_cpu = non_spec_mask_cpu & (query_lens_cpu == 1)
prefill_mask_cpu = non_spec_mask_cpu & (query_lens_cpu > 1)
num_spec_decodes = int(spec_sequence_masks_cpu.sum().item())
num_decodes = int(decode_mask_cpu.sum().item())
num_prefills = int(prefill_mask_cpu.sum().item())
num_decode_tokens = num_decodes
num_prefill_tokens = int(query_lens_cpu[prefill_mask_cpu].sum().item())
num_spec_decode_tokens = int(
query_lens_cpu[spec_sequence_masks_cpu].sum().item()
)
# Max prefill query length, CPU-side (no device-to-host sync) for the
# PLE dilated short-conv packing buffer.
max_prefill_query_len = (
int(query_lens_cpu[prefill_mask_cpu].max().item())
if num_prefills > 0
else 0
)
assert num_accepted_tokens is not None
non_spec_req_idx: torch.Tensor | None = None
if num_decodes == 0 and num_prefills == 0:
# All real requests are speculative; zero-length padding is trailing.
spec_token_indx = None
non_spec_token_indx = None
spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, 0]
num_accepted_tokens = num_accepted_tokens[:num_spec_decodes]
non_spec_state_indices_tensor = None
spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
non_spec_query_start_loc = None
else:
# Mixed batch: build a per-token group key consistent with the
# request grouping above (spec=0 | decode=1 | prefill=2) and a
# stable sort, so tokens of each request stay contiguous and in
# request order. This yields spec tokens first, then the non-spec
# [decode, prefill] tokens.
query_lens = torch.diff(query_start_loc)
spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0]
decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0]
prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0]
non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu))
spec_req_idx = async_tensor_h2d(
spec_req_idx_cpu, device=query_start_loc.device
)
non_spec_req_idx = async_tensor_h2d(
non_spec_req_idx_cpu, device=query_start_loc.device
)
decode_req_idx = async_tensor_h2d(
decode_req_idx_cpu, device=query_start_loc.device
)
req_group = torch.full(
(m.num_reqs,),
2,
dtype=torch.int64,
device=query_start_loc.device,
)
req_group.index_fill_(0, spec_req_idx, 0)
req_group.index_fill_(0, decode_req_idx, 1)
token_group = torch.repeat_interleave(
req_group, query_lens, output_size=int(query_start_loc_cpu[-1])
)
token_perm = torch.argsort(token_group, stable=True)
spec_token_indx = token_perm[:num_spec_decode_tokens]
non_spec_token_indx = token_perm[num_spec_decode_tokens:]
spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0]
num_accepted_tokens = num_accepted_tokens[spec_req_idx]
non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0]
spec_query_start_loc = torch.zeros(
num_spec_decodes + 1,
dtype=torch.int32,
device=query_start_loc.device,
)
torch.cumsum(query_lens[spec_req_idx], dim=0, out=spec_query_start_loc[1:])
non_spec_query_start_loc = torch.zeros(
num_decodes + num_prefills + 1,
dtype=torch.int32,
device=query_start_loc.device,
)
torch.cumsum(
query_lens[non_spec_req_idx],
dim=0,
out=non_spec_query_start_loc[1:],
)
# Compute the conv-state slots for the non-spec decode/prefill split.
if non_spec_state_indices_tensor is None:
state_indices_tensor = block_table_tensor[:0, 0]
else:
state_indices_tensor = non_spec_state_indices_tensor
# Build the regular decode/prefill state metadata inherited from the
# generic short-conv metadata contract.
query_start_loc_p = None
query_start_loc_d = None
state_indices_tensor_p = None
state_indices_tensor_d = None
if num_decodes > 0 or num_prefills > 0:
num_computed_tokens = m.compute_num_computed_tokens()
assert non_spec_req_idx is not None
num_computed_tokens = num_computed_tokens[non_spec_req_idx]
state_indices_tensor_d = state_indices_tensor[:num_decodes]
state_indices_tensor_p = state_indices_tensor[
num_decodes : num_decodes + num_prefills
]
if num_decodes > 0:
has_initial_states_d = num_computed_tokens[:num_decodes] > 0
assert non_spec_query_start_loc is not None
query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1]
if num_prefills > 0:
has_initial_states_p = (
num_computed_tokens[num_decodes : num_decodes + num_prefills] > 0
)
assert non_spec_query_start_loc is not None
query_start_loc_p = (
non_spec_query_start_loc[num_decodes:] - num_decode_tokens
)
# Prepare persistent tensors for CUDA graph capture and replay.
# ``m.num_actual_tokens`` is already padded by the model runner.
# Request-level buffers use ``m.num_reqs`` while token-level buffers
# use their independently bounded token count.
batch_size = m.num_reqs
if (
self.use_full_cuda_graph
and num_prefills == 0
and num_decodes == 0
and spec_sequence_masks is not None
and num_spec_decodes <= self.decode_cudagraph_max_bs
and num_spec_decode_tokens <= self.decode_cudagraph_max_tokens
):
assert spec_state_indices_tensor is not None
self.spec_state_indices_tensor[:num_spec_decodes].copy_(
spec_state_indices_tensor, non_blocking=True
)
spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID)
self.spec_sequence_masks[:batch_size].copy_(
spec_sequence_masks[:batch_size], non_blocking=True
)
spec_sequence_masks = self.spec_sequence_masks[:batch_size]
assert spec_query_start_loc is not None
self.spec_query_start_loc[: num_spec_decodes + 1].copy_(
spec_query_start_loc, non_blocking=True
)
spec_num_query_tokens = spec_query_start_loc[-1]
spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens)
assert num_accepted_tokens is not None
self.num_accepted_tokens[:num_spec_decodes].copy_(
num_accepted_tokens, non_blocking=True
)
num_accepted_tokens = self.num_accepted_tokens[:batch_size]
num_accepted_tokens[num_spec_decodes:].fill_(1)
return PleShortConvAttentionMetadata(
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
num_reqs=m.num_reqs,
num_spec_decodes=num_spec_decodes,
num_spec_decode_tokens=num_spec_decode_tokens,
num_actual_tokens=m.num_actual_tokens,
spec_query_len=self.num_spec + 1,
max_prefill_query_len=max_prefill_query_len,
query_start_loc=query_start_loc,
state_indices_tensor=state_indices_tensor,
has_initial_states_p=has_initial_states_p,
has_initial_states_d=has_initial_states_d,
non_spec_query_start_loc=non_spec_query_start_loc,
spec_query_start_loc=spec_query_start_loc,
spec_state_indices_tensor=spec_state_indices_tensor,
spec_sequence_masks=spec_sequence_masks,
spec_token_indx=spec_token_indx,
non_spec_token_indx=non_spec_token_indx,
num_accepted_tokens=num_accepted_tokens,
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
query_start_loc_p=query_start_loc_p,
query_start_loc_d=query_start_loc_d,
state_indices_tensor_p=state_indices_tensor_p,
state_indices_tensor_d=state_indices_tensor_d,
seq_lens=m.seq_lens,
)
def build_for_cudagraph_capture(
self, common_attn_metadata: CommonAttentionMetadata
) -> PleShortConvAttentionMetadata:
"""Build metadata for full CUDA graph capture.
Currently, only decode is supported for full CUDA graphs with
short-conv.
"""
m = common_attn_metadata
assert (
m.num_reqs <= self.decode_cudagraph_max_bs
and m.num_actual_tokens <= self.decode_cudagraph_max_tokens
), (
"ShortConv only supports decode-only full CUDAGraph capture. "
f"Make sure batch size ({m.num_reqs}) <= "
f"cudagraph capture size ({self.decode_cudagraph_max_bs}) and "
f"number of tokens ({m.num_actual_tokens}) <= "
f"token capture size ({self.decode_cudagraph_max_tokens})."
)
if self.use_spec_decode:
num_accepted_tokens = torch.diff(m.query_start_loc)
num_decode_draft_tokens_cpu = torch.diff(m.query_start_loc_cpu).sub_(1)
assert num_decode_draft_tokens_cpu.shape == num_accepted_tokens.shape
return self.build(
0,
m,
num_accepted_tokens=num_accepted_tokens,
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
)
return self.build(0, m)