WA for numerically unstable block_softmax (#104)

HabanaAI · Jul 17, 2024 · 34e2855 · 34e2855
1 parent f0e4a83
commit 34e2855
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/vllm/hpu/ops.py b/vllm/hpu/ops.py
@@ -42,10 +42,12 @@ def block2batch(tensor, block_mapping):
 
 
 def block_softmax(batch_size, attn, block_mapping):
+    attn.sub_(10.0)
     attn = attn.exp_()
     sums = attn.sum(dim=-1).unsqueeze(-1)
     sums = block2batch(sums, block_mapping)
     sums = batch2block(sums, block_mapping)
+    sums.add_(1.0e-12)
     attn.div_(sums)
     return attn