Possible support for weighted average loss calculation in FusedLinearCrossEntropy kernel

[ChenlongDeng]

🚀 The feature, motivation and pitch

Hi, great job on the FusedLinearCrossEntropy kernel! I’ve found it very efficient for model training. However, it seems to lack support for custom token-wise weighted loss calculation. For instance, users might input a 2D tensor containing individual loss weights for each token and calculate a weighted average loss based on those values. Does this version support that functionality, or is it something that might be added in the future?

Alternatives

No response

Additional context

No response