From 18f4c96bdb6a348ec7128afd319a98ec587bd5d1 Mon Sep 17 00:00:00 2001 From: maskyuanzh Date: Tue, 26 May 2026 10:40:42 +0000 Subject: [PATCH 1/4] Optimize grad_weight accumulation with addmm --- .../ops/fused_linear_cross_entropy.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 01f1b5658..96b8ed063 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -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 @@ -92,7 +94,7 @@ def fused_linear_cross_entropy_forward( ce_weight_sum = ce_weight.sum().item() if ce_weight.stride(-1) != 1: ce_weight = ce_weight.contiguous() - + IS_TORCH2P12 = Version(torch.__version__.split("+")[0]) >= Version("2.12.0") for chunk_id in range(num_chunks): start_idx = chunk_id * chunk_size end_idx = min((chunk_id + 1) * chunk_size, BT) @@ -209,7 +211,26 @@ 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 ( + IS_TORCH2P12 + 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.to(dtype=grad_logits_chunk.t().dtype), + out_dtype=torch.float32, + out=grad_weight, + ) + else: + torch.addmm( + grad_weight, + grad_logits_chunk.t().to(grad_weight.dtype), + _input_chunk.to(grad_weight.dtype), + out=grad_weight, + ) if bias is not None and input_requires_grad: torch.add( From d4977713b2877726540de3b26db63f5c4d1b4ac7 Mon Sep 17 00:00:00 2001 From: maskyuanzh Date: Wed, 27 May 2026 07:22:12 +0000 Subject: [PATCH 2/4] fix fused linear CE addmm fallback dtype handling --- src/liger_kernel/ops/fused_linear_cross_entropy.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 96b8ed063..80ee0404e 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -14,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]) +_SUPPORTS_ADDM_MIXED_PRECISION_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") def fused_linear_cross_entropy_forward( @@ -41,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 @@ -94,7 +95,7 @@ def fused_linear_cross_entropy_forward( ce_weight_sum = ce_weight.sum().item() if ce_weight.stride(-1) != 1: ce_weight = ce_weight.contiguous() - IS_TORCH2P12 = Version(torch.__version__.split("+")[0]) >= Version("2.12.0") + for chunk_id in range(num_chunks): start_idx = chunk_id * chunk_size end_idx = min((chunk_id + 1) * chunk_size, BT) @@ -212,7 +213,7 @@ def fused_linear_cross_entropy_forward( if grad_weight is not None and input_requires_grad: if ( - IS_TORCH2P12 + _SUPPORTS_ADDM_MIXED_PRECISION_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) @@ -225,12 +226,7 @@ def fused_linear_cross_entropy_forward( out=grad_weight, ) else: - torch.addmm( - grad_weight, - grad_logits_chunk.t().to(grad_weight.dtype), - _input_chunk.to(grad_weight.dtype), - out=grad_weight, - ) + grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float() if bias is not None and input_requires_grad: torch.add( From ec69ffb07088f5e1917cfc51e980be7ae48dccec Mon Sep 17 00:00:00 2001 From: maskyuanzh Date: Wed, 24 Jun 2026 05:20:34 +0000 Subject: [PATCH 3/4] Remove redundant input chunk cast --- src/liger_kernel/ops/fused_linear_cross_entropy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 80ee0404e..8f2919ceb 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -221,7 +221,7 @@ def fused_linear_cross_entropy_forward( torch.addmm( grad_weight, grad_logits_chunk.t(), - _input_chunk.to(dtype=grad_logits_chunk.t().dtype), + _input_chunk, out_dtype=torch.float32, out=grad_weight, ) From 4c21cc71414a4263bb5e763101eff3bfcd3f0964 Mon Sep 17 00:00:00 2001 From: maskyuanzh Date: Wed, 1 Jul 2026 02:01:13 +0000 Subject: [PATCH 4/4] Rename addmm out_dtype support constant --- src/liger_kernel/ops/fused_linear_cross_entropy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/liger_kernel/ops/fused_linear_cross_entropy.py b/src/liger_kernel/ops/fused_linear_cross_entropy.py index 8f2919ceb..6685f733a 100644 --- a/src/liger_kernel/ops/fused_linear_cross_entropy.py +++ b/src/liger_kernel/ops/fused_linear_cross_entropy.py @@ -15,7 +15,7 @@ # 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]) -_SUPPORTS_ADDM_MIXED_PRECISION_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") +_ADDMM_SUPPORTS_OUT_DTYPE = _TORCH_VERSION >= Version("2.8.0") def fused_linear_cross_entropy_forward( @@ -213,7 +213,7 @@ def fused_linear_cross_entropy_forward( if grad_weight is not None and input_requires_grad: if ( - _SUPPORTS_ADDM_MIXED_PRECISION_OUT_DTYPE + _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)