vllm.v1.attention.backends.mla.triton_mla ¶

class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
    _cudagraph_support: ClassVar[AttentionCGSupport] = (
        AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
    )
    # Non-causal DSpark block is flattened to one decode row per query token in
    # forward_mqa, so no intra-block causal masking is required.
    supports_non_causal_multi_token_decode: ClassVar[bool] = True

    @classmethod
    def get_cudagraph_support(
        cls,
        vllm_config: VllmConfig,
        kv_cache_spec: KVCacheSpec,
    ) -> AttentionCGSupport:
        """Report UNIFORM_BATCH where a non-causal multi-token block is served.

        ``_cudagraph_support`` is a class constant, so serving the DSpark
        draft's (1 + num_spec) block through the decode path reports
        UNIFORM_SINGLE_TOKEN_DECODE and, because the engine takes the minimum
        over all attention groups, downgrades the *whole* engine off full
        cudagraphs. ``forward_mqa`` flattens that block with
        ``repeat_interleave`` on a Python int and performs no device->host
        sync, so it does satisfy the UNIFORM_BATCH contract.

        ``non_causal_multi_token_decode`` is a KV-cache-group property, not a
        per-layer one: ``MLAAttentionSpec.merge`` ORs it over every layer in
        the group, so a group holding both a draft and its target reports it
        for both. That is the same predicate ``__init__`` below already uses to
        raise ``reorder_batch_threshold``, so the two stay consistent, but it
        does mean this lifts a causal target sharing the draft's KV cache group
        as well.
        """
        if getattr(kv_cache_spec, "non_causal_multi_token_decode", False):
            return AttentionCGSupport.UNIFORM_BATCH
        return cls._cudagraph_support

    def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
        # DCP local sequence lengths are not advanced between draft steps.
        self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
        # Only the non-causal DSpark draft group serves multi-token blocks via
        # the decode path; raise its reorder threshold to the spec block length
        # so full-cudagraph capture admits it. Causal usage stays single-token.
        if getattr(self, "non_causal_multi_token_decode", False):
            self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
        self._reserve_attn_logits_workspace()

    def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None:
        pass

    def _reserve_attn_logits_workspace(self) -> None:
        """Pre-size the shared workspace for the decode split-KV attn logits.

        Reserving at the worst case (max_model_len -> max num_kv_splits,
        max_num_seqs decode tokens) before warmup/cudagraph capture means the
        per-call ``get_simultaneous`` in ``forward_mqa`` never has to grow the
        buffer at runtime (which would raise once the workspace is locked).
        """
        if not is_workspace_manager_initialized():
            return
        # Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
        B = self.vllm_config.scheduler_config.max_num_seqs
        # Non-causal DSpark draft flattens each request's block to query_len
        # decode rows; cover max_num_seqs * block_len rows.
        if getattr(self, "non_causal_multi_token_decode", False):
            B *= self.reorder_batch_threshold
        # DCP all-gathers the query heads before forward_mqa.
        q_num_heads = self.num_heads * self.dcp_world_size
        max_splits = _compute_num_kv_splits(
            self.model_config.max_model_len,
            current_platform.num_compute_units(),
        )
        lse_dim = self.mla_dims.kv_lora_rank + 1
        current_workspace_manager().get_simultaneous(
            ((B, q_num_heads, max_splits, lse_dim), torch.float32),
        )