[Benchmark] kl_div kernel and test #615
Conversation
stack-info: PR: #615, branch: Sibylau/stack/1
Sibylau
force-pushed
the
Sibylau/stack/1
branch
from
dc48a33 to
384353d
Compare
meta-cla
bot
added
the
CLA Signed
| elif reduction == "mean": | ||
| final_loss = torch.sum(loss) / (BT * V) | ||
| else: # reduction == "none" | ||
| final_loss = loss |
There was a problem hiding this comment.
should we test all reductions in the unit tests / main function? or maybe simplify it to just one case that's being used by tritonbench?
There was a problem hiding this comment.
tritonbench only tests the reduction = "batchmean" case: https://github.com/meta-pytorch/tritonbench/blob/main/tritonbench/operators/kl_div/operator.py#L28
while liger_kernel implements all cases. I mirrored liger_kernel implementation: https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/kl_div.py#L150
There was a problem hiding this comment.
would you suggest I add all unit tests for all reductions, or only implement reduction = "batchmean" ?
There was a problem hiding this comment.
I see, let's just keep it as-is now.
|
codegen test passes |
stack-info: PR: #615, branch: Sibylau/stack/1
yf225
force-pushed
the
Sibylau/stack/1
branch
from
3878bf5 to
ad5be78
Compare
|
I took the liberty to rebase the PR. Will merge it after the CI clears. |
Stacked PRs:
[Benchmark] kl_div kernel and test