[Optimization] CrossEntropyLoss

[StrongSpoon]

PR Category

Operator

Type of Change

Performance Optimization

Description

apply fusion on celoss kernel and sum kernel of weight.
avoid memory copy from device to host by disabling context saving when input tensor doesn't require gradient.
achieve 90% relative performance on average compared to aten (only 13% before)
performance is better in benchmark with cudagraph or on RTX 3090

A100-benchmark without cudagraph

A100-benchmark with cudagraph

Issue

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance