diff --git a/vllm/model_executor/layers/rotary_embedding/base.py b/vllm/model_executor/layers/rotary_embedding/base.py index 1374334b2c..41291a9b56 100644 --- a/vllm/model_executor/layers/rotary_embedding/base.py +++ b/vllm/model_executor/layers/rotary_embedding/base.py @@ -62,6 +62,17 @@ class RotaryEmbeddingBase(CustomOp): self.cos_sin_cache: torch.Tensor self.register_buffer("cos_sin_cache", cache, persistent=False) + # Reuse a precomputed bf16 cache for the AITER compile path. + if self.use_aiter and cache.dtype != torch.bfloat16: + self.cos_sin_cache_bf16: torch.Tensor | None + self.register_buffer( + "cos_sin_cache_bf16", + cache.to(torch.bfloat16), + persistent=False, + ) + else: + self.cos_sin_cache_bf16 = None + self.apply_rotary_emb = ApplyRotaryEmb( is_neox_style=self.is_neox_style, ) @@ -101,6 +112,16 @@ class RotaryEmbeddingBase(CustomOp): ): return cos_sin_cache + # Reuse precomputed bf16 cache in the AITER compile path. + if ( + self.use_aiter + and torch.compiler.is_compiling() + and query.dtype == torch.bfloat16 + ): + cache_bf16 = getattr(self, "cos_sin_cache_bf16", None) + if cache_bf16 is not None and cache_bf16.device == query.device: + return cache_bf16 + cos_sin_cache = cos_sin_cache.to(query.device, dtype=query.dtype) # Avoid mutating buffers during torch.compile (cudagraph) tracing. if torch.compiler.is_compiling():