[Misc] Update type annotation for rotary embedding base (#18914)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2025-05-30 10:17:01 +08:00
parent d54af615d5
commit 1aa2f81b43
4 changed files with 23 additions and 26 deletions
--- a/benchmarks/kernels/benchmark_rope.py
+++ b/benchmarks/kernels/benchmark_rope.py
@@ -22,7 +22,7 @@ def benchmark_rope_kernels_multi_lora(
    seed: int,
    device: str,
    max_position: int = 8192,
-    base: int = 10000,
+    base: float = 10000,
 ) -> None:
    current_platform.seed_everything(seed)
    torch.set_default_device(device)