[BugFix][NPU] Fix Fish Speech S2 Pro NPU support (#7546)

Signed-off-by: Lei Ke <1141466880@qq.com>

tests/attention/test_fish_kvcache_attn.py

@@ -3,6 +3,7 @@
 
 import pytest
 import torch
+from vllm.platforms import current_platform
 
 from vllm_omni.attention import fish_kvcache_attn, fish_kvcache_backend
 
@@ -113,6 +114,15 @@ def test_fish_kvcache_enabled_by_default(monkeypatch):
     assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
 
 
+def test_fish_kvcache_fast_path_is_unavailable_off_cuda(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(
+        fish_kvcache_attn, "_triton_backend", lambda: pytest.fail("Triton backend must not load off CUDA")
+    )
+
+    assert not fish_kvcache_attn.is_available()
+
+
 @pytest.mark.parametrize("value", ["", "1", "true", "yes", "on", "required"])
 def test_fish_kvcache_enabled_values_are_consistent(monkeypatch, value):
     monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)

tests/config/test_config_factory.py

@@ -2688,6 +2688,13 @@ class TestPlatformOverrides:
             # Explicit null clears the inherited single-GPU 2 GiB CUDA cap.
             assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None
 
+    def test_fish_speech_npu_uses_ascend_kv_block_size(self):
+        deploy_path = Path(get_deploy_config_path("fish_qwen3_omni.yaml"))
+
+        deploy = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu")
+
+        assert deploy.stages[0].engine_extras["block_size"] == 128
+
     def test_npu_overrides(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():

vllm_omni/attention/fish_kvcache_attn.py

@@ -24,6 +24,14 @@ def _triton_backend() -> Any:
 
 
 def is_available() -> bool:
+    # The Fish kvcache fast path is CUDA-only: it depends on FlashAttention's
+    # use_cascade metadata (absent from Ascend's metadata) and a Triton CUDA
+    # kernel. On other platforms report unavailable so the native attention
+    # backend takes over instead of crashing on attn_metadata.use_cascade.
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_cuda():
+        return False
     return _triton_backend().is_available()
 
 

vllm_omni/deploy/fish_qwen3_omni.yaml

@@ -64,6 +64,15 @@ stages:
       repetition_penalty: 1.0
 
 platforms:
+  npu:
+    # Ascend paged attention only accepts 128-token cache blocks
+    # (AscendAttentionBackend.get_supported_kernel_block_sizes() == [128]),
+    # while vLLM's CacheConfig default is 16. Without this override stage 0
+    # aborts during KV-cache init with:
+    #   ValueError: No common block size for 16.
+    stages:
+      - stage_id: 0
+        block_size: 128
   xpu:
     stages:
       # Verified on Intel Arc B60.

vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py

@@ -32,6 +32,7 @@ from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.platforms import current_omni_platform
 from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig
@@ -195,7 +196,11 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         self.mtp_hidden_size = int(self.text_config.hidden_size)
         self.talker_mtp_output_key = ("codes", "audio")
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")}
-        self.talker_mtp_graph_safe = True
+        # NPU (Ascend) ACL graphs cannot capture the Fast-AR random-sampling ops
+        # (torch.nonzero / torch.multinomial used for top-p), so keep the talker
+        # MTP eager on NPU to fall back to native (non-graph) execution. Mirrors
+        # the 310P patch for Qwen3-TTS (talker_mtp_graph_safe=False).
+        self.talker_mtp_graph_safe = not current_omni_platform.is_npu()
 
         # Qwen3 transformer backbone.
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))