@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1aa2f81b43

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 07:15 UTC

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

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed May 30, 2025 at 10:17 UTC 1aa2f81b43b286ec9fa66439c151d18f20a662ec

4 files changed +23 -26

benchmarks/kernels/benchmark_rope.py

+1 -1

@@ -22,7 +22,7 @@ def benchmark_rope_kernels_multi_lora(

22 seed: int,

23 device: str,

24 max_position: int = 8192,

25 - base: int = 10000,

25 + base: float = 10000,

26 ) -> None:

27 current_platform.seed_everything(seed)

28 torch.set_default_device(device)

tests/kernels/core/test_pos_encoding.py

+3 -3

@@ -70,7 +70,7 @@ def test_rotary_embedding(

70 device: str,

71 use_key: bool,

72 max_position: int = 8192,

73 - base: int = 10000,

73 + base: float = 10000,

74 ) -> None:

75 if rotary_dim is None:

76 rotary_dim = head_size

@@ -135,7 +135,7 @@ def test_batched_rotary_embedding(

135 device: str,

136 use_key: bool,

137 max_position: int = 8192,

138 - base: int = 10000,

138 + base: float = 10000,

139 ) -> None:

140 current_platform.seed_everything(seed)

141 torch.set_default_device(device)

@@ -203,7 +203,7 @@ def test_batched_rotary_embedding_multi_lora(

203 device: str,

204 use_key: bool,

205 max_position: int = 8192,

206 - base: int = 10000,

206 + base: float = 10000,

207 ) -> None:

208 current_platform.seed_everything(seed)

209 torch.set_default_device(device)

vllm/model_executor/layers/rotary_embedding.py

+17 -17

@@ -96,7 +96,7 @@ class RotaryEmbedding(CustomOp):

96 head_size: int,

97 rotary_dim: int,

98 max_position_embeddings: int,

99 - base: int,

99 + base: float,

100 is_neox_style: bool,

101 dtype: torch.dtype,

102 ) -> None:

@@ -113,7 +113,7 @@ class RotaryEmbedding(CustomOp):

113 self.cos_sin_cache: torch.Tensor

114 self.register_buffer("cos_sin_cache", cache, persistent=False)

115

116 - def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:

116 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

117 """Compute the inverse frequency."""

118 # NOTE(woosuk): To exactly match the HF implementation, we need to

119 # use CPU to compute the cache and then move it to GPU. However, we

@@ -404,7 +404,7 @@ class LinearScalingRotaryEmbedding(RotaryEmbedding):

404 head_size: int,

405 rotary_dim: int,

406 max_position_embeddings: int,

407 - base: int,

407 + base: float,

408 is_neox_style: bool,

409 scaling_factors: Union[list[float], float],

410 dtype: torch.dtype,

@@ -464,7 +464,7 @@ class NTKScalingRotaryEmbedding(RotaryEmbedding):

464 head_size: int,

465 rotary_dim: int,

466 max_position_embeddings: int,

467 - base: int,

467 + base: float,

468 is_neox_style: bool,

469 scaling_factor: float,

470 dtype: torch.dtype,

@@ -474,7 +474,7 @@ class NTKScalingRotaryEmbedding(RotaryEmbedding):

474 super().__init__(head_size, rotary_dim, max_position_embeddings, base,

475 is_neox_style, dtype)

476

477 - def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:

477 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

478 base = self.base * (self.scaling_factor if self.mixed_b is None else 1)

479 inv_freq = super()._compute_inv_freq(base)

480

@@ -501,7 +501,7 @@ class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):

501 head_size: int,

502 rotary_dim: int,

503 max_position_embeddings: int,

504 - base: int,

504 + base: float,

505 is_neox_style: bool,

506 scaling_factor: float,

507 dtype: torch.dtype,

@@ -582,7 +582,7 @@ class YaRNScalingRotaryEmbedding(RotaryEmbedding):

582 head_size: int,

583 rotary_dim: int,

584 max_position_embeddings: int,

585 - base: int,

585 + base: float,

586 is_neox_style: bool,

587 scaling_factor: float,

588 dtype: torch.dtype,

@@ -644,7 +644,7 @@ class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):

644 rotary_dim: int,

645 max_position_embeddings: int,

646 original_max_position_embeddings: int,

647 - base: int,

647 + base: float,

648 is_neox_style: bool,

649 dtype: torch.dtype,

650 short_factor: list[float],

@@ -769,7 +769,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):

769 head_size: int,

770 rotary_dim: int,

771 max_position_embeddings: int,

772 - base: int,

772 + base: float,

773 is_neox_style: bool,

774 scaling_factor: float,

775 dtype: torch.dtype,

@@ -877,7 +877,7 @@ class Llama3RotaryEmbedding(RotaryEmbedding):

877 head_size: int,

878 rotary_dim: int,

879 max_position_embeddings: int,

880 - base: int,

880 + base: float,

881 is_neox_style: bool,

882 dtype: torch.dtype,

883 scaling_factor: float,

@@ -892,7 +892,7 @@ class Llama3RotaryEmbedding(RotaryEmbedding):

892 super().__init__(head_size, rotary_dim, max_position_embeddings, base,

893 is_neox_style, dtype)

894

895 - def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:

895 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

896 inv_freqs = super()._compute_inv_freq(base)

897 low_freq_wavelen = self.orig_max_position / self.low_freq_factor

898 high_freq_wavelen = self.orig_max_position / self.high_freq_factor

@@ -923,14 +923,14 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):

923 head_size: int,

924 rotary_dim: int,

925 max_position_embeddings: int,

926 - base: int,

926 + base: float,

927 is_neox_style: bool,

928 dtype: torch.dtype,

929 ):

930 super().__init__(head_size, rotary_dim, max_position_embeddings, base,

931 is_neox_style, dtype)

932

933 - def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:

933 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

934 inv_freqs = super()._compute_inv_freq(base)

935 inv_freqs = inv_freqs[:(self.rotary_dim // 2)]

936 return inv_freqs

@@ -989,7 +989,7 @@ class MRotaryEmbedding(RotaryEmbedding):

989 head_size: int,

990 rotary_dim: int,

991 max_position_embeddings: int,

992 - base: int,

992 + base: float,

993 is_neox_style: bool,

994 dtype: torch.dtype,

995 mrope_section: Optional[list[int]] = None,

@@ -1529,7 +1529,7 @@ class DualChunkRotaryEmbedding(CustomOp):

1529 head_size: int,

1530 rotary_dim: int,

1531 max_position_embeddings: int,

1532 - base: int,

1532 + base: float,

1533 is_neox_style: bool,

1534 dtype: torch.dtype,

1535 chunk_size: int,

@@ -1558,7 +1558,7 @@ class DualChunkRotaryEmbedding(CustomOp):

1558 q_inter_cache,

1559 persistent=False)

1560

1561 - def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:

1561 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

1562 """Compute the inverse frequency."""

1563 # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`.

1564 # However, we use `torch.arange(..., dtype=torch.float)` instead to

@@ -1705,7 +1705,7 @@ def get_rope(

1705 head_size: int,

1706 rotary_dim: int,

1707 max_position: int,

1708 - base: int,

1708 + base: float,

1709 is_neox_style: bool = True,

1710 rope_scaling: Optional[dict[str, Any]] = None,

1711 dtype: Optional[torch.dtype] = None,

vllm/model_executor/models/minimax_text_01.py

+2 -5

@@ -141,7 +141,7 @@ class MiniMaxText01RotaryEmbedding(CustomOp):

141 head_size: int,

142 rotary_dim: int,

143 max_position: int,

144 - base: int,

144 + base: float,

145 is_neox_style: bool,

146 cache_dtype: torch.dtype,

147 ) -> None:

@@ -155,10 +155,7 @@ class MiniMaxText01RotaryEmbedding(CustomOp):

155 cache = self._compute_cos_sin_cache().to(cache_dtype)

156 self.register_buffer("cos_sin_cache", cache, persistent=False)

157

158 - def _compute_inv_freq(

159 - self,

160 - base: Union[int, float],

161 - ) -> torch.Tensor:

158 + def _compute_inv_freq(self, base: float) -> torch.Tensor:

159 """Compute the inverse frequency."""

160 inv_freq = 1.0 / (base**(torch.arange(

161 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))