Skip to content

vllm_omni.platforms.npu.models.cosyvoice2_dit_attn

NPU patches for CosyVoice2 / Token2Wav DiT attention.

CosyVoice2's DiT Attention.forward builds a key-padding style mask (B, 1, 1, S) via mask.unsqueeze(1) and feeds it to F.scaled_dot_product_attention. On Ascend that call is routed to npu_fusion_attention / aclnnFlashAttentionScore, which only accepts mask shapes [B,N,S,S] / [B,1,S,S] / [1,1,S,S] / [S,S] and fails with error 161001 (tiling / parameter invalid) for [B,1,1,S].

This module: 1. Expands DiT attention masks to [B, 1, S, S] before SDPA. 2. Provides a MATH-backend SDPA context so inference can avoid the fused FA kernel entirely when the platform still incorrectly routes SDPA.

logger module-attribute 

logger = init_logger(__name__)

apply_cosyvoice2_dit_attn_npu_patch 

apply_cosyvoice2_dit_attn_npu_patch() -> None

Monkey-patch CosyVoice2 DiT Attention for Ascend FA mask constraints.

npu_math_sdpa_context 

npu_math_sdpa_context(
    *, require_available: bool = False
) -> Iterator[None]

Force SDPA MATH backend so Ascend does not call fused FA.