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, 22:00 UTC
Reduce memory usage for granite_speech. (#42933)
Signed-off-by: Yihuki <wangbovbvb@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Yihuki committed
May 25, 2026 at 14:12 UTC
5c1aec3dc0600cb816a0389e55ef1f1c17893380
1 file changed
+1
-3
vllm/model_executor/models/granite_speech.py
+1
-3
@@ -389,10 +389,8 @@ class GraniteSpeechConformerAttention(nn.Module):
389
# shaw's relative positional embedding
390
dist = attention_dists.to(hidden_states.device)
391
rel_pos_emb = self.rel_pos_emb(dist)
392
- rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] + list(rel_pos_emb.shape))
392
pos_attn = (
394
- torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded, dim=-1)
395
- * self.scale
393
+ torch.einsum("bnhid,ijd->bnhij", query_states, rel_pos_emb) * self.scale
394
)
395
396
if remainder > 0: