From 7c6369b0c762c867db9bae7a15cff01e948ba6ff Mon Sep 17 00:00:00 2001
From: xthomaswang <xthomaswang@gmail.com>
Date: Thu, 23 Apr 2026 00:55:18 -0400
Subject: [PATCH] [Relax][NN] Use int64 for RoPE apply flag

---
 python/tvm/relax/frontend/nn/llm/position_embedding.py | 2 +-
 src/runtime/vm/paged_kv_cache.cc                       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index dec80c50c270..cec2ba65dcdb 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -529,7 +529,7 @@ def fused_rope(  # pylint: disable=too-many-locals
         var_q: T.handle,
         var_k: T.handle,
         var_v: T.handle,
-        apply_rope: T.int32,
+        apply_rope: T.int64,
     ):
         T.func_attr(
             {
diff --git a/src/runtime/vm/paged_kv_cache.cc b/src/runtime/vm/paged_kv_cache.cc
index 36f7697237e2..696f80e0c86c 100644
--- a/src/runtime/vm/paged_kv_cache.cc
+++ b/src/runtime/vm/paged_kv_cache.cc
@@ -1337,7 +1337,7 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj {
     }
     if (!rope_ext_factors_.defined()) {
       f_split_rotary_(qkv_data_view, q_rope_position_map_view_, q_data, k_data, v_data,
-                      static_cast<int>(rope_mode_ == RoPEMode::kNormal));
+                      static_cast<int64_t>(rope_mode_ == RoPEMode::kNormal));
     } else {
       f_split_rotary_(qkv_data_view, q_rope_position_map_view_, q_data, k_data, v_data,
                       rope_ext_factors_.value());