search

pytorch-labs / tritonbench

Tritonbench is a collection of PyTorch custom operators with example inputs to measure their performance.
BSD 3-Clause "New" or "Revised" License
20 stars 3 forks source link

Add custom ops jsd and fused_linear_jsd #16

Closed FindHao closed 3 weeks ago

FindHao commented 1 month ago

Migrated from https://github.com/pytorch/benchmark/pull/2518

Add custom ops jsd and fused_linear_jsd

Test Plan:

% python run.py --op jsd,fused_linear_jsd  --num-inputs 1
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:06<00:00,  6.94s/it]
  x_val    torch_jsd-latency    liger_jsd-latency    inductor_jsd-latency
-------  -------------------  -------------------  ----------------------
      0               2.1768             0.461984                0.154944
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:04<00:00,  4.69s/it]
  x_val    torch_lm_head_jsd-latency    liger_lm_head_jsd-latency    inductor_lm_head_jsd-latency
-------  ---------------------------  ---------------------------  ------------------------------
      0                      73.6553                      362.348                         66.4232
facebook-github-bot commented 1 month ago

@FindHao has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.