Skip to content

Commit

Permalink
Remove workaround for one_hot in eager/compile (#632)
Browse files Browse the repository at this point in the history
As now one_hot operator has implementation for eager and compile mode
the workaround is not needed any longer.
  • Loading branch information
anko-intel authored Dec 16, 2024
1 parent 449a89d commit d2128b4
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions vllm/worker/hpu_model_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,11 +296,11 @@ def _set_block_mapping(self, metadata, batch_size, device, dtype):
attn_bias = (torch.zeros_like(mask, dtype=dtype).masked_fill_(
mask, -math.inf))

if not is_fake_hpu() and htorch.utils.internal.is_lazy():
if not is_fake_hpu():
block_mapping = torch.nn.functional.one_hot(metadata.block_groups,
num_classes=batch_size)
else:
# Unfortunately one_hot on CPU/torch.compile mode/eager mode
# Unfortunately one_hot on CPU
# doesn't handle out of bounds classes so we need to convert
# all negative values to 0 (block_mapping) or bs (block_groups)
block_groups = metadata.block_groups.to(torch.long)
Expand Down

0 comments on commit d2128b4

Please sign in to comment.