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

    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),
        )