Skip to content

perf: optimize grad_weight accumulation with addmm #1239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vaibhavjindal merged 5 commits into from
Jul 1, 2026
+19 −2
Merged

perf: optimize grad_weight accumulation with addmm #1239

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
18f4c96
Optimize grad_weight accumulation with addmm
May 26, 2026
d497771
fix fused linear CE addmm fallback dtype handling
May 27, 2026
ec69ffb
Remove redundant input chunk cast
Jun 24, 2026
69a1575
Merge branch 'main' into fix-fused-linear-ce-addmm
vaibhavjindal Jun 30, 2026
4c21cc7
Rename addmm out_dtype support constant
Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions src/liger_kernel/ops/fused_linear_cross_entropy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import torch
import triton

from packaging.version import Version

from liger_kernel.ops.cross_entropy import liger_cross_entropy_kernel
from liger_kernel.ops.utils import amp_custom_bwd
from liger_kernel.ops.utils import amp_custom_fwd
Expand All @@ -12,6 +14,8 @@
# However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling
# The optimal maximum block size depends on your hardware, your kernel, and your dtype
MAX_FUSED_SIZE = 2048 if infer_device() == "npu" else 65536 // 2
_TORCH_VERSION = Version(torch.__version__.split("+")[0])
_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0")


def fused_linear_cross_entropy_forward(
Expand Down Expand Up @@ -39,7 +43,6 @@ def fused_linear_cross_entropy_forward(
f"return_predicted_tokens must be True or False. Got: {return_predicted_tokens}"
)
device = _input.device

input_requires_grad = _input.requires_grad

# inputs have shape: BT x H
Expand Down Expand Up @@ -209,7 +212,21 @@ def fused_linear_cross_entropy_forward(
grad_input[start_idx:end_idx] = grad_logits_chunk @ weight

if grad_weight is not None and input_requires_grad:
grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()
if (
_ADDMM_SUPPORTS_OUT_DTYPE
and grad_weight.device.type == "cuda"
and grad_weight.dtype == torch.float32
and grad_logits_chunk.t().dtype in (torch.float16, torch.bfloat16)
):
torch.addmm(
grad_weight,
grad_logits_chunk.t(),
_input_chunk,
out_dtype=torch.float32,

@tyler-romero tyler-romero May 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think technically out_dtype is available earlier than 2.12, I just didnt do the work to track down which version it was introduced in. I think I remember it existing in 2.10 as well.

@maskyuanzh maskyuanzh May 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. The exact version wasn’t carefully verified here. After checking the PyTorch docs and source tags, I found that out_dtype was added to torch.addmm in PyTorch 2.8.0 for fp16/bf16 CUDA inputs with fp32 output accumulation. I’ll lower the version guard from 2.12.0 to 2.8.0.

out=grad_weight,
)
else:
grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()

if bias is not None and input_requires_grad:
torch.add(
Expand Down
Loading