swap sin cos, add deepseek_v3 attention test

Author: nihui

src/layer/rotaryembed.cpp

@@ -21,8 +21,8 @@ int RotaryEmbed::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
     // assert bottom_blobs.size() == 3
 
     const Mat& bottom_blob = bottom_blobs[0];
-    const Mat& sin_cache = bottom_blobs[1];
-    const Mat& cos_cache = bottom_blobs[2];
+    const Mat& cos_cache = bottom_blobs[1];
+    const Mat& sin_cache = bottom_blobs[2];
 
     const int embed_dim = bottom_blob.w;
     const int seqlen = bottom_blob.h;
@@ -44,16 +44,16 @@ int RotaryEmbed::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
             if (interleaved)
             {
                 const float* ptr = head.row(i);
-                const float* sin_ptr = sin_cache.row(i);
                 const float* cos_ptr = cos_cache.row(i);
+                const float* sin_ptr = sin_cache.row(i);
                 float* outptr = out_head.row(i);
 
                 for (int j = 0; j < embed_dim / 2; j++)
                 {
                     const float x1 = ptr[0];
                     const float x2 = ptr[1];
-                    const float sin_val = *sin_ptr++;
                     const float cos_val = *cos_ptr++;
+                    const float sin_val = *sin_ptr++;
                     outptr[0] = x1 * cos_val - x2 * sin_val;
                     outptr[1] = x1 * sin_val + x2 * cos_val;
                     ptr += 2;
@@ -73,8 +73,8 @@ int RotaryEmbed::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>&
                 {
                     const float x1 = *ptr1++;
                     const float x2 = *ptr2++;
-                    const float sin_val = *sin_ptr++;
                     const float cos_val = *cos_ptr++;
+                    const float sin_val = *sin_ptr++;
                     *outptr1++ = x1 * cos_val - x2 * sin_val;
                     *outptr2++ = x1 * sin_val + x2 * cos_val;
                 }

tools/pnnx/src/CMakeLists.txt

@@ -430,6 +430,7 @@ set(pnnx_pass_ncnn_SRCS
     pass_ncnn/eliminate_output.cpp
     pass_ncnn/expand_expression.cpp
     pass_ncnn/fuse_convert_shufflechannel_slice.cpp
+    pass_ncnn/fuse_convert_rotaryembed.cpp
     pass_ncnn/insert_split.cpp
     pass_ncnn/chain_multi_output.cpp
     pass_ncnn/solve_batch_index.cpp

tools/pnnx/src/pass_ncnn.cpp

@@ -23,6 +23,7 @@
 #include "pass_ncnn/eliminate_output.h"
 #include "pass_ncnn/expand_expression.h"
 #include "pass_ncnn/fuse_convert_shufflechannel_slice.h"
+#include "pass_ncnn/fuse_convert_rotaryembed.h"
 #include "pass_ncnn/insert_split.h"
 #include "pass_ncnn/chain_multi_output.h"
 #include "pass_ncnn/solve_batch_index.h"
@@ -79,6 +80,8 @@ void pass_ncnn(Graph& g, const std::vector<std::string>& module_operators)
 
     attribute_unpooling(g);
 
+    ncnn::fuse_convert_rotaryembed(g);
+
     ncnn::expand_expression(g);
 
     ncnn::chain_multi_output(g);

tools/pnnx/src/pass_ncnn/fuse_convert_rotaryembed.cpp

@@ -0,0 +1,82 @@
+// Copyright 2025 Tencent
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "fuse_convert_rotaryembed.h"
+
+#include "pass_level2.h"
+
+namespace pnnx {
+
+namespace ncnn {
+
+class fuse_rotaryembed_pass_interleaved : public GraphRewriterPass
+{
+public:
+    const char* match_pattern_graph() const
+    {
+        return R"PNNXIR(7767517
+11 11
+pnnx.Input              input_0     0 1 input
+pnnx.Input              input_1     0 1 cos_cache
+pnnx.Input              input_2     0 1 sin_cache
+Tensor.reshape          op_0        1 1 input 22 shape=(%batch,%num_heads,%seqlen,%embed_dim_half,2)
+torch.transpose         op_1        1 1 22 23 dim0=%interleave_dim0 dim1=%interleave_dim1
+Tensor.reshape          op_2        1 1 23 24 shape=(%batch,%num_heads,%seqlen,%embed_dim)
+torch.tensor_split      op_3        1 2 24 28 29 dim=%split_dim indices=(%embed_dim_half)
+pnnx.Expression         op_4        1 1 29 30 expr=neg(@0)
+torch.cat               op_5        2 1 30 28 31 dim=%cat_dim
+pnnx.Expression         op_6        4 1 24 cos_cache 31 sin_cache out expr=add(mul(@0,@1),mul(@2,@3))
+pnnx.Output             output      1 0 out
+)PNNXIR";
+    }
+
+    const char* type_str() const
+    {
+        return "RotaryEmbed";
+    }
+
+    const char* name_str() const
+    {
+        return "rope";
+    }
+
+    bool match(const std::map<std::string, Parameter>& captured_params) const
+    {
+        const int embed_dim_half = captured_params.at("embed_dim_half").i;
+        const int embed_dim = captured_params.at("embed_dim").i;
+        if (embed_dim != embed_dim_half * 2)
+            return false;
+
+        const int interleave_dim0 = captured_params.at("interleave_dim0").i;
+        const int interleave_dim1 = captured_params.at("interleave_dim1").i;
+        if (!((interleave_dim0 == 4 && interleave_dim1 == 3) || (interleave_dim0 == 3 && interleave_dim1 == 4)))
+            return false;
+
+        const int split_dim = captured_params.at("split_dim").i;
+        if (split_dim != 3 && split_dim != -1)
+            return false;
+
+        const int cat_dim = captured_params.at("cat_dim").i;
+        if (cat_dim != 3 && cat_dim != -1)
+            return false;
+
+        return true;
+    }
+
+    void write(Operator* op, const std::map<std::string, Parameter>& /*captured_params*/) const
+    {
+        op->params["0"] = 1; // interleaved
+    }
+};
+
+void fuse_convert_rotaryembed(Graph& graph)
+{
+    fuse_rotaryembed_pass_interleaved a;
+    int opindex = 0;
+
+    pnnx_graph_rewrite(graph, &a, opindex);
+}
+
+} // namespace ncnn
+
+} // namespace pnnx

tools/pnnx/src/pass_ncnn/fuse_convert_rotaryembed.h

@@ -0,0 +1,14 @@
+// Copyright 2025 Tencent
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "ir.h"
+
+namespace pnnx {
+
+namespace ncnn {
+
+void fuse_convert_rotaryembed(Graph& graph);
+
+} // namespace ncnn
+
+} // namespace pnnx

tools/pnnx/tests/CMakeLists.txt

@@ -383,6 +383,7 @@ pnnx_add_test(transformers_clip_attention)
 pnnx_add_test(transformers_chinese_clip_attention)
 pnnx_add_test(transformers_ctrl_attention)
 pnnx_add_test(transformers_deberta_attention)
+pnnx_add_test(transformers_deepseek_v3_attention)
 pnnx_add_test(transformers_distilbert_attention)
 pnnx_add_test(transformers_electra_attention)
 pnnx_add_test(transformers_flaubert_attention)

tools/pnnx/tests/test_transformers_deepseek_v3_attention.py

@@ -0,0 +1,60 @@
+# Copyright 2025 Tencent
+# SPDX-License-Identifier: BSD-3-Clause
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from packaging import version
+
+if version.parse(torch.__version__) < version.parse('2.1'):
+    exit(0)
+
+from transformers import DeepseekV3Config
+from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3Attention, DeepseekV3RotaryEmbedding
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        config = DeepseekV3Config(hidden_size=192, num_attention_heads=16, num_key_value_heads=16, q_lora_rank=64, kv_lora_rank=128, attn_implementation='sdpa')
+        self.rotary_emb = DeepseekV3RotaryEmbedding(config)
+        self.attn0 = DeepseekV3Attention(config, layer_idx=1)
+
+    def forward(self, x, mask0):
+        batch_size = x.size(0)
+        seq_length = x.size(1)
+        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+        position_embeddings = self.rotary_emb(x, position_ids)
+        out0 = self.attn0(x, position_embeddings=position_embeddings, attention_mask=mask0, head_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, output_attentions=True)
+        return out0[0]
+
+def test():
+    net = Model()
+    net.eval()
+
+    torch.manual_seed(0)
+    x = torch.rand(3, 16, 192)
+
+    mask0 = torch.rand(3, 1, 16, 16)
+
+    a = net(x, mask0)
+
+    # export torchscript
+    mod = torch.jit.trace(net, (x, mask0))
+    mod.save("test_transformers_deepseek_v3_attention.pt")
+
+    # torchscript to pnnx
+    import os
+    os.system("../src/pnnx test_transformers_deepseek_v3_attention.pt inputshape=[3,16,192],[3,1,16,16] fp16=0")
+
+    # pnnx inference
+    import test_transformers_deepseek_v3_attention_pnnx
+    b = test_transformers_deepseek_v3_attention_pnnx.test_inference()
+
+    return torch.allclose(a, b, 1e-4, 1e-4)
+
+if __name__ == "__main__":
+    if test():
+        exit(0)
+    else:
+        exit(1)