Skip to content

Commit

Permalink
Add HSTU apply_rope and batched_multihead_index_sum
Browse files Browse the repository at this point in the history
Summary: Add apply_rope operator from HSTU

Reviewed By: int3

Differential Revision: D61150155

fbshipit-source-id: 0317c45efc87e00d208ec2ce8749278314bbafad
  • Loading branch information
xuzhao9 authored and facebook-github-bot committed Aug 21, 2024
1 parent 57de51f commit 5a121b6
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions torchbenchmark/operators/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,10 @@ def load_opbench_by_name(op_name: str):
)
opbench_list = list(opbench_list)
if not opbench_list:
raise RuntimeError(f"{op_name} is not found in the core model list.")
raise RuntimeError(f"{op_name} is not found in the Tritonbench operator list.")
assert (
len(opbench_list) == 1
), f"Found more than one models {opbench_list} with the exact name: {op_name}"
), f"Found more than one operators {opbench_list} matching the required name: {op_name}"
op_name = opbench_list[0]
op_pkg = (
op_name
Expand Down

0 comments on commit 5a121b6

Please sign in to comment.