fix(ascend): fall back for unsupported RoPE dtype

Author: hooray03

tests/unit_tests/dispatch/test_ascend_rotary.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+"""Tests for the Ascend rotary embedding implementation."""
+
+import torch
+
+from vllm_fl.dispatch.backends.reference.impl.rotary import rotary_embedding_torch
+from vllm_fl.dispatch.backends.vendor.ascend.impl.rotary import rotary_embedding_ascend
+
+
+def test_ascend_rotary_falls_back_for_float32():
+    num_tokens = 4
+    num_heads = 2
+    head_size = 8
+
+    query = torch.randn(num_tokens, num_heads, head_size)
+    key = torch.randn_like(query)
+    positions = torch.arange(num_tokens)
+    frequencies = torch.randn(16, head_size // 2)
+    cos = frequencies.cos()
+    sin = frequencies.sin()
+
+    actual_query, actual_key = rotary_embedding_ascend(
+        None, query, key, cos, sin, positions, inplace=False
+    )
+    expected_query, expected_key = rotary_embedding_torch(
+        None, query, key, cos, sin, positions, inplace=False
+    )
+
+    torch.testing.assert_close(actual_query, expected_query)
+    torch.testing.assert_close(actual_key, expected_key)

vllm_fl/dispatch/backends/vendor/ascend/impl/rotary.py

@@ -36,6 +36,26 @@ def rotary_embedding_ascend(
     Returns:
         Tuple of (embedded_query, embedded_key)
     """
+    # The ATB rotary kernel used by torch-npu does not support FP32 inputs.
+    # Launch failures are reported asynchronously and otherwise surface in a
+    # later, unrelated operator. Keep the optimized path for model dtypes and
+    # use the reference implementation for unsupported dtypes.
+    if query.dtype not in (torch.float16, torch.bfloat16):
+        from vllm_fl.dispatch.backends.reference.impl.rotary import (
+            rotary_embedding_torch,
+        )
+
+        return rotary_embedding_torch(
+            obj,
+            query,
+            key,
+            cos,
+            sin,
+            position_ids,
+            rotary_interleaved=rotary_interleaved,
+            inplace=inplace,
+        )
+
     import torch_npu
 
     # query/key shape: [num_tokens, num_heads, rotary_dim]