WA for numerically unstable block_softmax (#104)

HabanaAI · Aug 7, 2024 · 530f70c · 530f70c
1 parent 810fadb
commit 530f70c
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