[RELAX] Fix rotary embedding buffer size calculation (#18102)

harrywhoo · web-flow · commit 910aeafb8432 · 2025-07-01T00:04:07.000-04:00
* Change head_dim//2 to rotary_dim//2 in LongRope scaling
* Fixes buffer size when rotary_dim differs from head_dim
diff --git 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 fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
         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(
             {

Original file line number	Diff line number	Diff line change
`@@ -493,7 +493,7 @@ def fused_rope_longrope_scaling( # pylint: disable=too-many-locals`
`493`	`493`	`var_q: T.handle,`
`494`	`494`	`var_k: T.handle,`
`495`	`495`	`var_v: T.handle,`
`496`		`- ext_factors: T.Buffer((head_dim // 2,), "float32"), # type: ignore`
	`496`	`+ ext_factors: T.Buffer((rotary_dim // 2,), "float32"), # type: ignore`
`497`	`497`	`):`
`498`	`498`	`T.func_attr(`
`499`	`499`	`{`