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 2, 18:15 UTC
[ROCm][GPT-OSS] Avoid repeated compile-time `cos_sin_cache.to(bf16)` casts in rotary path (#42833)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
akii96 committed
May 27, 2026 at 11:22 UTC
de12f5ca0b8bcec55b6fa29174900360ef31cc0e
1 file changed
+21
vllm/model_executor/layers/rotary_embedding/base.py
+21
@@ -62,6 +62,17 @@ class RotaryEmbeddingBase(CustomOp):
62
self.cos_sin_cache: torch.Tensor
63
self.register_buffer("cos_sin_cache", cache, persistent=False)
64
65
+ # Reuse a precomputed bf16 cache for the AITER compile path.
66
+ if self.use_aiter and cache.dtype != torch.bfloat16:
67
+ self.cos_sin_cache_bf16: torch.Tensor | None
68
+ self.register_buffer(
69
+ "cos_sin_cache_bf16",
70
+ cache.to(torch.bfloat16),
71
+ persistent=False,
72
+ )
73
+ else:
74
+ self.cos_sin_cache_bf16 = None
75
+
76
self.apply_rotary_emb = ApplyRotaryEmb(
77
is_neox_style=self.is_neox_style,
78
)
@@ -101,6 +112,16 @@ class RotaryEmbeddingBase(CustomOp):
112
):
113
return cos_sin_cache
114
115
+ # Reuse precomputed bf16 cache in the AITER compile path.
116
+ if (
117
+ self.use_aiter
118
+ and torch.compiler.is_compiling()
119
+ and query.dtype == torch.bfloat16
120
+ ):
121
+ cache_bf16 = getattr(self, "cos_sin_cache_bf16", None)
122
+ if cache_bf16 is not None and cache_bf16.device == query.device:
123
+ return cache_bf16
124
+
125
cos_sin_cache = cos_sin_cache.to(query.device, dtype=query.dtype)
126
# Avoid mutating buffers during torch.compile (cudagraph) tracing.
127
if torch.compiler.is_compiling():