(tvm) branch main updated: [RELAX] Fix rotary embedding buffer size calculation (#18102)

[ruihangl]

This is an automated email from the ASF dual-hosted git repository.

ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 910aeafb84 [RELAX] Fix rotary embedding buffer size calculation 
(#18102)
910aeafb84 is described below

commit 910aeafb8432db7ea07336a0deb72d3a123d38b7
Author: Yuehan Hu <harryhu...@gmail.com>
AuthorDate: Tue Jul 1 00:04:07 2025 -0400

    [RELAX] Fix rotary embedding buffer size calculation (#18102)
    
    * Change head_dim//2 to rotary_dim//2 in LongRope scaling
    * Fixes buffer size when rotary_dim differs from head_dim
---
 python/tvm/relax/frontend/nn/llm/position_embedding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py 
b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index fc82148be1..1a1659b29e 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -493,7 +493,7 @@ def llama_rope_with_position_map(  # pylint: 
disable=too-many-arguments
         var_q: T.handle,
         var_k: T.handle,
         var_v: T.handle,
-        ext_factors: T.Buffer((head_dim // 2,), "float32"),  # type: ignore
+        ext_factors: T.Buffer((rotary_dim // 2,), "float32"),  # type: ignore
     ):
         T.func_attr(
             {