From 83cd58309bce3df43a152e4227bd854a55c40423 Mon Sep 17 00:00:00 2001 From: vasqu Date: Wed, 15 Jul 2026 19:54:18 +0200 Subject: [PATCH 1/3] add device context --- .../torch-ext/liger_kernels/rms_norm.py | 189 +++++++++--------- .../torch-ext/liger_kernels/utils.py | 11 + 2 files changed, 108 insertions(+), 92 deletions(-) diff --git a/liger-kernels/torch-ext/liger_kernels/rms_norm.py b/liger-kernels/torch-ext/liger_kernels/rms_norm.py index 52e103cb..0a76a5a5 100644 --- a/liger-kernels/torch-ext/liger_kernels/rms_norm.py +++ b/liger-kernels/torch-ext/liger_kernels/rms_norm.py @@ -24,6 +24,8 @@ from .utils import set_large_grf_mode from .utils import torch_to_triton_dtype from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -438,47 +440,49 @@ def rms_norm_forward(X, W, eps, offset, casting_mode, row_mode): kernel_args = {} if X.device.type == "xpu": set_large_grf_mode(kernel_args) - if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode: - _rms_norm_forward_kernel[(n_rows,)]( - Y, - Y.stride(0), - X, - X.stride(0), - W, - W.stride(0) if elementwise_affine else 0, - RSTD, - RSTD.stride(0), - n_cols, - eps, - offset, - casting_mode, - elementwise_affine=elementwise_affine, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - **kernel_args, # XPU-specific optimization - ) - else: - BLOCK_ROW = 16 - kernel_args["BLOCK_ROW"] = BLOCK_ROW - _block_rms_norm_forward_kernel[(triton.cdiv(n_rows, BLOCK_ROW),)]( - Y, - Y.stride(0), - X, - X.stride(0), - W, - W.stride(0) if elementwise_affine else 0, - RSTD, - RSTD.stride(0), - n_rows, - n_cols, - eps, - offset, - casting_mode, - elementwise_affine=elementwise_affine, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - **kernel_args, # XPU-specific optimization - ) + + with device_context(X.device): + if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode: + _rms_norm_forward_kernel[(n_rows,)]( + Y, + Y.stride(0), + X, + X.stride(0), + W, + W.stride(0) if elementwise_affine else 0, + RSTD, + RSTD.stride(0), + n_cols, + eps, + offset, + casting_mode, + elementwise_affine=elementwise_affine, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + **kernel_args, # XPU-specific optimization + ) + else: + BLOCK_ROW = 16 + kernel_args["BLOCK_ROW"] = BLOCK_ROW + _block_rms_norm_forward_kernel[(triton.cdiv(n_rows, BLOCK_ROW),)]( + Y, + Y.stride(0), + X, + X.stride(0), + W, + W.stride(0) if elementwise_affine else 0, + RSTD, + RSTD.stride(0), + n_rows, + n_cols, + eps, + offset, + casting_mode, + elementwise_affine=elementwise_affine, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + **kernel_args, # XPU-specific optimization + ) return Y.view(*shape), X, RSTD, BLOCK_SIZE, num_warps, casting_mode @@ -519,57 +523,58 @@ def rms_norm_backward(dY, X, W, RSTD, offset, casting_mode, BLOCK_SIZE, num_warp if X.device.type == "xpu": set_large_grf_mode(kernel_args) - if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode: - _rms_norm_backward_kernel[grid]( - dY, - dY.stride(0), - dX, - dX.stride(0), - X, - X.stride(0), - torch_to_triton_dtype[X.dtype], - W, - W.stride(0) if elementwise_affine else 0, - RSTD, - RSTD.stride(0), - _dW, - _dW.stride(0) if elementwise_affine else 0, - n_rows, - n_cols, - offset, - rows_per_program, - casting_mode, - elementwise_affine=elementwise_affine, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - **kernel_args, # XPU-specific optimization - ) - else: - BLOCK_ROW = 16 - kernel_args["BLOCK_ROW"] = BLOCK_ROW - _block_rms_norm_backward_kernel[grid]( - dY, - dY.stride(0), - dX, - dX.stride(0), - X, - X.stride(0), - torch_to_triton_dtype[X.dtype], - W, - W.stride(0) if elementwise_affine else 0, - RSTD, - RSTD.stride(0), - _dW, - _dW.stride(0) if elementwise_affine else 0, - n_rows, - n_cols, - offset, - casting_mode, - elementwise_affine=elementwise_affine, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - **kernel_args, # XPU-specific optimization - ) + with device_context(X.device): + if BLOCK_SIZE > 256 or n_rows < 4096 * 8 or row_mode: + _rms_norm_backward_kernel[grid]( + dY, + dY.stride(0), + dX, + dX.stride(0), + X, + X.stride(0), + torch_to_triton_dtype[X.dtype], + W, + W.stride(0) if elementwise_affine else 0, + RSTD, + RSTD.stride(0), + _dW, + _dW.stride(0) if elementwise_affine else 0, + n_rows, + n_cols, + offset, + rows_per_program, + casting_mode, + elementwise_affine=elementwise_affine, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + **kernel_args, # XPU-specific optimization + ) + else: + BLOCK_ROW = 16 + kernel_args["BLOCK_ROW"] = BLOCK_ROW + _block_rms_norm_backward_kernel[grid]( + dY, + dY.stride(0), + dX, + dX.stride(0), + X, + X.stride(0), + torch_to_triton_dtype[X.dtype], + W, + W.stride(0) if elementwise_affine else 0, + RSTD, + RSTD.stride(0), + _dW, + _dW.stride(0) if elementwise_affine else 0, + n_rows, + n_cols, + offset, + casting_mode, + elementwise_affine=elementwise_affine, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + **kernel_args, # XPU-specific optimization + ) dX = dX.view(*shape) if elementwise_affine: diff --git a/liger-kernels/torch-ext/liger_kernels/utils.py b/liger-kernels/torch-ext/liger_kernels/utils.py index 7de26149..53fffdf0 100644 --- a/liger-kernels/torch-ext/liger_kernels/utils.py +++ b/liger-kernels/torch-ext/liger_kernels/utils.py @@ -21,6 +21,7 @@ import triton.language as tl from packaging.version import Version +from contextlib import contextmanager def is_npu_available() -> bool: @@ -174,3 +175,13 @@ def set_large_grf_mode(kernel_args: dict): else: # API was changed in https://github.com/intel/intel-xpu-backend-for-triton/pull/5430 kernel_args["grf_mode"] = "large" + +@contextmanager +def device_context(device: torch.device): + """Context manager that sets the active device for any backend (cuda, xpu, etc.).""" + backend = getattr(torch, device.type, None) + if backend is not None and hasattr(backend, "device"): + with backend.device(device): + yield + else: + yield From 9e9db35482a2cfaf3a5550681820114dcccc4c6f Mon Sep 17 00:00:00 2001 From: vasqu Date: Thu, 16 Jul 2026 14:27:16 +0200 Subject: [PATCH 2/3] try same fix for gelgu/swiglu --- .../torch-ext/liger_kernels/geglu.py | 40 +++++++++-------- .../torch-ext/liger_kernels/swiglu.py | 43 ++++++++++--------- 2 files changed, 45 insertions(+), 38 deletions(-) diff --git a/liger-kernels/torch-ext/liger_kernels/geglu.py b/liger-kernels/torch-ext/liger_kernels/geglu.py index c2cf1995..cab9d8cd 100644 --- a/liger-kernels/torch-ext/liger_kernels/geglu.py +++ b/liger-kernels/torch-ext/liger_kernels/geglu.py @@ -8,6 +8,8 @@ from .utils import compare_version from .utils import ensure_contiguous from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -94,15 +96,16 @@ def geglu_forward(a, b): BLOCK_SIZE, num_warps = calculate_settings(n_cols) - _geglu_tanh_forward_kernel[(n_rows,)]( - a, - b, - c, - c.stride(-2), - n_cols=n_cols, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) + with device_context(a.device): + _geglu_tanh_forward_kernel[(n_rows,)]( + a, + b, + c, + c.stride(-2), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) return a, b, c.view(*ori_shape) @@ -114,15 +117,16 @@ def geglu_backward(a, b, dc): BLOCK_SIZE, num_warps = calculate_settings(n_cols) - _geglu_tanh_backward_kernel[(n_rows,)]( - dc, - a, - b, - dc.stride(-2), - n_cols=n_cols, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) + with device_context(a.device): + _geglu_tanh_backward_kernel[(n_rows,)]( + dc, + a, + b, + dc.stride(-2), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) return a.view(*ori_shape), b.view(*ori_shape) diff --git a/liger-kernels/torch-ext/liger_kernels/swiglu.py b/liger-kernels/torch-ext/liger_kernels/swiglu.py index d77b90fe..9776e44a 100644 --- a/liger-kernels/torch-ext/liger_kernels/swiglu.py +++ b/liger-kernels/torch-ext/liger_kernels/swiglu.py @@ -4,6 +4,7 @@ from .utils import calculate_settings from .utils import ensure_contiguous +from .utils import device_context @triton.jit @@ -73,16 +74,17 @@ def swiglu_forward(a, b, gate_multiplier: float = 1.0): BLOCK_SIZE, num_warps = calculate_settings(n_cols) - _swiglu_forward_kernel[(n_rows,)]( - a, - b, - c, - c.stride(-2), - float(gate_multiplier), - n_cols=n_cols, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) + with device_context(a.device): + _swiglu_forward_kernel[(n_rows,)]( + a, + b, + c, + c.stride(-2), + float(gate_multiplier), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) return a, b, c.view(*ori_shape) @@ -94,16 +96,17 @@ def swiglu_backward(a, b, dc, gate_multiplier: float = 1.0): BLOCK_SIZE, num_warps = calculate_settings(n_cols) - _swiglu_backward_kernel[(n_rows,)]( - dc, - a, - b, - dc.stride(-2), - float(gate_multiplier), - n_cols=n_cols, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - ) + with device_context(a.device): + _swiglu_backward_kernel[(n_rows,)]( + dc, + a, + b, + dc.stride(-2), + float(gate_multiplier), + n_cols=n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + ) return a.view(*ori_shape), b.view(*ori_shape) From 95f9074ecceaaddb2753cc7300a297d4eaed5a63 Mon Sep 17 00:00:00 2001 From: vasqu Date: Fri, 17 Jul 2026 17:49:51 +0200 Subject: [PATCH 3/3] apply device guard everywhere around triton launches --- .../torch-ext/liger_kernels/cross_entropy.py | 104 ++++++------- liger-kernels/torch-ext/liger_kernels/dyt.py | 28 ++-- .../fused_linear_cross_entropy.py | 141 +++++++++--------- .../torch-ext/liger_kernels/group_norm.py | 82 +++++----- liger-kernels/torch-ext/liger_kernels/jsd.py | 36 ++--- .../torch-ext/liger_kernels/kl_div.py | 53 +++---- .../torch-ext/liger_kernels/layer_norm.py | 90 +++++------ .../torch-ext/liger_kernels/qwen2vl_mrope.py | 78 +++++----- liger-kernels/torch-ext/liger_kernels/rope.py | 90 +++++------ liger-kernels/torch-ext/liger_kernels/tvd.py | 38 ++--- 10 files changed, 387 insertions(+), 353 deletions(-) diff --git a/liger-kernels/torch-ext/liger_kernels/cross_entropy.py b/liger-kernels/torch-ext/liger_kernels/cross_entropy.py index 52933cb5..fdcdba9e 100644 --- a/liger-kernels/torch-ext/liger_kernels/cross_entropy.py +++ b/liger-kernels/torch-ext/liger_kernels/cross_entropy.py @@ -11,6 +11,8 @@ from .utils import is_hip from .utils import infer_device from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -372,44 +374,45 @@ def cross_entropy_forward( if target.stride(-1) != 1: target = target.contiguous() - # Here we use a trick to store X_ptr gradient in X_ptr so we can save memory - liger_cross_entropy_kernel[(n_rows,)]( - X_ptr=_input, - X_stride=_input.stride(-2), - Y_ptr=target, - Y_stride=target.stride(-1), # always 1 - weight_ptr=weight, # dummy if None - loss_ptr=loss_1d, - z_loss_ptr=z_loss_1d, - loss_stride=loss_1d.stride(-1), # always 1 - token_accuracy_ptr=token_accuracy_1d, - token_accuracy_stride=token_accuracy_1d.stride(-1) - if return_token_accuracy - else 0, # always 1 if accuracy is enabled - predicted_tokens_ptr=predicted_tokens_1d, - predicted_tokens_stride=predicted_tokens_1d.stride(-1) - if return_predicted_tokens - else 0, # always 1 if predicted tokens is enabled - n_cols=V, - n_non_ignore=n_non_ignore, - sum_non_ignore_weight=sum_non_ignore_weight, - ignore_index=ignore_index, - weight_sum=weight_sum, - lse_square_scale=lse_square_scale, - label_smoothing=label_smoothing, - reduction=reduction, - softcap=softcap, - RETURN_Z_LOSS=return_z_loss, - RETURN_TOKEN_ACCURACY=return_token_accuracy, - RETURN_PREDICTED_TOKENS=return_predicted_tokens, - BLOCK_SIZE=BLOCK_SIZE, - HAS_WEIGHT=True if weight is not None else False, - HAS_SOFTCAPPING=True if softcap is not None else False, - HAS_GRADIENTS=_input.requires_grad, - # TODO: 32 seems to give the best performance - # Performance is quite sensitive to num_warps - num_warps=32 if not is_hip() else 16, - ) + with device_context(_input.device): + # Here we use a trick to store X_ptr gradient in X_ptr so we can save memory + liger_cross_entropy_kernel[(n_rows,)]( + X_ptr=_input, + X_stride=_input.stride(-2), + Y_ptr=target, + Y_stride=target.stride(-1), # always 1 + weight_ptr=weight, # dummy if None + loss_ptr=loss_1d, + z_loss_ptr=z_loss_1d, + loss_stride=loss_1d.stride(-1), # always 1 + token_accuracy_ptr=token_accuracy_1d, + token_accuracy_stride=token_accuracy_1d.stride(-1) + if return_token_accuracy + else 0, # always 1 if accuracy is enabled + predicted_tokens_ptr=predicted_tokens_1d, + predicted_tokens_stride=predicted_tokens_1d.stride(-1) + if return_predicted_tokens + else 0, # always 1 if predicted tokens is enabled + n_cols=V, + n_non_ignore=n_non_ignore, + sum_non_ignore_weight=sum_non_ignore_weight, + ignore_index=ignore_index, + weight_sum=weight_sum, + lse_square_scale=lse_square_scale, + label_smoothing=label_smoothing, + reduction=reduction, + softcap=softcap, + RETURN_Z_LOSS=return_z_loss, + RETURN_TOKEN_ACCURACY=return_token_accuracy, + RETURN_PREDICTED_TOKENS=return_predicted_tokens, + BLOCK_SIZE=BLOCK_SIZE, + HAS_WEIGHT=True if weight is not None else False, + HAS_SOFTCAPPING=True if softcap is not None else False, + HAS_GRADIENTS=_input.requires_grad, + # TODO: 32 seems to give the best performance + # Performance is quite sensitive to num_warps + num_warps=32 if not is_hip() else 16, + ) if reduction == "none": loss = loss_1d @@ -437,18 +440,19 @@ def cross_entropy_backward(_input, grad_output): # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. else: - BT, V = _input.shape - n_rows = BT - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) - - element_mul_kernel[(n_rows,)]( - _input, - _input.stride(-2), - grad_output, - V, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32 if not is_hip() else 16, - ) + with device_context(_input.device): + BT, V = _input.shape + n_rows = BT + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(V)) + + element_mul_kernel[(n_rows,)]( + _input, + _input.stride(-2), + grad_output, + V, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=32 if not is_hip() else 16, + ) return _input diff --git a/liger-kernels/torch-ext/liger_kernels/dyt.py b/liger-kernels/torch-ext/liger_kernels/dyt.py index fd596b2e..3ab8606e 100644 --- a/liger-kernels/torch-ext/liger_kernels/dyt.py +++ b/liger-kernels/torch-ext/liger_kernels/dyt.py @@ -9,6 +9,8 @@ from .utils import get_npu_core_count from .utils import infer_device from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -107,16 +109,17 @@ def liger_dyt_fwd(x, alpha, gamma, beta): y = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M) - _dyt_fwd_kernel[grid]( - x, - y, - alpha, - gamma, - beta, - HAVE_BETA, - N, - ) + with device_context(x.device): + grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), M) + _dyt_fwd_kernel[grid]( + x, + y, + alpha, + gamma, + beta, + HAVE_BETA, + N, + ) return y.view(input_shape) @@ -139,8 +142,9 @@ def liger_dyt_bwd(dy, x, alpha, gamma, beta): db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None dx = torch.empty_like(dy) - grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS) - _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N) + with device_context(x.device): + grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), NUM_SMS) + _dyt_bwd_kernel[grid](dy, dx, da, dg, db, x, alpha, gamma, HAVE_BETA, M, N) if HAVE_BETA: db = db.sum(0).to(x.dtype) dg = dg.sum(0).to(gamma.dtype) diff --git a/liger-kernels/torch-ext/liger_kernels/fused_linear_cross_entropy.py b/liger-kernels/torch-ext/liger_kernels/fused_linear_cross_entropy.py index e2197824..33a435d1 100644 --- a/liger-kernels/torch-ext/liger_kernels/fused_linear_cross_entropy.py +++ b/liger-kernels/torch-ext/liger_kernels/fused_linear_cross_entropy.py @@ -7,6 +7,7 @@ from .utils import element_mul_kernel from .utils import is_hip from .utils import infer_device +from .utils import device_context # The hard limit of TRITON_MAX_TENSOR_NUMEL is 1048576 https://github.com/triton-lang/triton/blob/ba42a5c68fd0505f8c42f4202d53be0f8d9a5fe0/python/triton/language/core.py#L19 # However, setting limit as 65536 as in LayerNorm tutorial is faster because of less register spilling @@ -147,42 +148,43 @@ def fused_linear_cross_entropy_forward( logits_chunk = logits_chunk.contiguous() target_chunk = target_chunk.contiguous() - # Here we calculate the gradient of logits_chunk in place so we can save memory. - liger_cross_entropy_kernel[(n_rows,)]( - X_ptr=logits_chunk, - X_stride=logits_chunk.stride(-2), - Y_ptr=target_chunk, - Y_stride=target_chunk.stride(-1), # always 1 - weight_ptr=ce_weight, - loss_ptr=loss_1d_slice, - z_loss_ptr=z_loss_1d_slice, - loss_stride=loss_1d_slice.stride(-1), # always 1 - token_accuracy_ptr=token_accuracy_1d_slice, - token_accuracy_stride=token_accuracy_1d_slice.stride(-1) - if return_token_accuracy - else 0, # always 1 if accuracy is enabled - predicted_tokens_ptr=predicted_tokens_1d_slice, - predicted_tokens_stride=predicted_tokens_1d_slice.stride(-1) - if return_predicted_tokens - else 0, # always 1 if predicted tokens is enabled - n_cols=V, - n_non_ignore=total_n_non_ignore, - sum_non_ignore_weight=total_sum_non_ignore_ce_weight, - weight_sum=ce_weight_sum, - ignore_index=ignore_index, - lse_square_scale=lse_square_scale, - label_smoothing=label_smoothing, - reduction=reduction, - softcap=softcap, - RETURN_Z_LOSS=return_z_loss, - RETURN_TOKEN_ACCURACY=return_token_accuracy, - RETURN_PREDICTED_TOKENS=return_predicted_tokens, - HAS_WEIGHT=True if ce_weight is not None else False, - HAS_SOFTCAPPING=True if softcap is not None else False, - HAS_GRADIENTS=input_requires_grad, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32 if not is_hip() else 16, - ) + with device_context(device): + # Here we calculate the gradient of logits_chunk in place so we can save memory. + liger_cross_entropy_kernel[(n_rows,)]( + X_ptr=logits_chunk, + X_stride=logits_chunk.stride(-2), + Y_ptr=target_chunk, + Y_stride=target_chunk.stride(-1), # always 1 + weight_ptr=ce_weight, + loss_ptr=loss_1d_slice, + z_loss_ptr=z_loss_1d_slice, + loss_stride=loss_1d_slice.stride(-1), # always 1 + token_accuracy_ptr=token_accuracy_1d_slice, + token_accuracy_stride=token_accuracy_1d_slice.stride(-1) + if return_token_accuracy + else 0, # always 1 if accuracy is enabled + predicted_tokens_ptr=predicted_tokens_1d_slice, + predicted_tokens_stride=predicted_tokens_1d_slice.stride(-1) + if return_predicted_tokens + else 0, # always 1 if predicted tokens is enabled + n_cols=V, + n_non_ignore=total_n_non_ignore, + sum_non_ignore_weight=total_sum_non_ignore_ce_weight, + weight_sum=ce_weight_sum, + ignore_index=ignore_index, + lse_square_scale=lse_square_scale, + label_smoothing=label_smoothing, + reduction=reduction, + softcap=softcap, + RETURN_Z_LOSS=return_z_loss, + RETURN_TOKEN_ACCURACY=return_token_accuracy, + RETURN_PREDICTED_TOKENS=return_predicted_tokens, + HAS_WEIGHT=True if ce_weight is not None else False, + HAS_SOFTCAPPING=True if softcap is not None else False, + HAS_GRADIENTS=input_requires_grad, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=32 if not is_hip() else 16, + ) # Apply token scaling if requested if use_token_scaling: @@ -247,47 +249,48 @@ def fused_linear_cross_entropy_forward( def fused_linear_cross_entropy_backward(grad_output, grad_input, grad_weight, grad_bias): # If cross entropy is the last layer, grad_output is 1.0. Skip the mul to save time if not torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)): - # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place - # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. - BT, H = grad_input.shape - n_rows = BT - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) - - element_mul_kernel[(n_rows,)]( - grad_input, - grad_input.stride(-2), - grad_output, - H, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32 if not is_hip() else 16, - ) - - # handle grad_weight - if grad_weight is not None: - V, H = grad_weight.shape - n_rows = V + with device_context(grad_input.device): + # We use a Triton kernel instead of a PyTorch operation because modifying inputs in-place + # for gradient storage and backward multiple times causes anomalies with PyTorch but not with Triton. + BT, H = grad_input.shape + n_rows = BT + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(H)) element_mul_kernel[(n_rows,)]( - grad_weight, - grad_weight.stride(-2), + grad_input, + grad_input.stride(-2), grad_output, H, BLOCK_SIZE=BLOCK_SIZE, num_warps=32 if not is_hip() else 16, ) - if grad_bias is not None: - V = grad_bias.shape[0] - n_rows = V - - element_mul_kernel[(n_rows,)]( - grad_bias, - grad_bias.stride(-1), - grad_output, - 1, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=32 if not is_hip() else 16, - ) + # handle grad_weight + if grad_weight is not None: + V, H = grad_weight.shape + n_rows = V + + element_mul_kernel[(n_rows,)]( + grad_weight, + grad_weight.stride(-2), + grad_output, + H, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=32 if not is_hip() else 16, + ) + + if grad_bias is not None: + V = grad_bias.shape[0] + n_rows = V + + element_mul_kernel[(n_rows,)]( + grad_bias, + grad_bias.stride(-1), + grad_output, + 1, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=32 if not is_hip() else 16, + ) return grad_input, grad_weight, grad_bias diff --git a/liger-kernels/torch-ext/liger_kernels/group_norm.py b/liger-kernels/torch-ext/liger_kernels/group_norm.py index 51c264bf..7b85363b 100644 --- a/liger-kernels/torch-ext/liger_kernels/group_norm.py +++ b/liger-kernels/torch-ext/liger_kernels/group_norm.py @@ -8,6 +8,8 @@ from .utils import ensure_contiguous from .utils import infer_device from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -215,26 +217,27 @@ def group_norm_forward(X, num_channels, num_groups, W, B, eps): Mean = torch.zeros((batch_size, num_groups), dtype=X.dtype, device=X.device) RSTD = torch.zeros((batch_size, num_groups), dtype=X.dtype, device=X.device) - _group_norm_forward_kernel[(batch_size, num_groups)]( - Y, - Y.stride(0), - Y.stride(1), - X, - X.stride(0), - X.stride(1), - Mean, - Mean.stride(0), - Mean.stride(1), - RSTD, - RSTD.stride(0), - RSTD.stride(1), - W, - B, - hidden_size, - channels_per_group, - eps, - BLOCK_SIZE=BLOCK_SIZE, - ) + with device_context(X.device): + _group_norm_forward_kernel[(batch_size, num_groups)]( + Y, + Y.stride(0), + Y.stride(1), + X, + X.stride(0), + X.stride(1), + Mean, + Mean.stride(0), + Mean.stride(1), + RSTD, + RSTD.stride(0), + RSTD.stride(1), + W, + B, + hidden_size, + channels_per_group, + eps, + BLOCK_SIZE=BLOCK_SIZE, + ) # Return tensors in the original shape return Y.view(*shape), X.view(*shape), Mean, RSTD, BLOCK_SIZE @@ -254,25 +257,26 @@ def group_norm_backward(dY, X, W, B, Mean, RSTD, num_channels, num_groups): DB = torch.zeros((num_channels), dtype=B.dtype, device=B.device) triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16 - BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(hidden_size)) - _group_norm_backward_kernel[(batch_size, num_groups)]( - X, - X.stride(0), - X.stride(1), - W, - Mean, - Mean.stride(0), - Mean.stride(1), - RSTD, - DX, - DW, - DB, - dY, - hidden_size, - channels_per_group, - BLOCK_SIZE=BLOCK_SIZE, - dtype=triton_dtype, - ) + with device_context(X.device): + BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(hidden_size)) + _group_norm_backward_kernel[(batch_size, num_groups)]( + X, + X.stride(0), + X.stride(1), + W, + Mean, + Mean.stride(0), + Mean.stride(1), + RSTD, + DX, + DW, + DB, + dY, + hidden_size, + channels_per_group, + BLOCK_SIZE=BLOCK_SIZE, + dtype=triton_dtype, + ) # Return tensors in the original shape return DX.view(*shape), DW, DB diff --git a/liger-kernels/torch-ext/liger_kernels/jsd.py b/liger-kernels/torch-ext/liger_kernels/jsd.py index 024ac908..fc8939d6 100644 --- a/liger-kernels/torch-ext/liger_kernels/jsd.py +++ b/liger-kernels/torch-ext/liger_kernels/jsd.py @@ -6,6 +6,7 @@ from .utils import ensure_contiguous from .utils import infer_device +from .utils import device_context @triton.jit @@ -109,23 +110,24 @@ def jsd_forward(_input, target, shift_labels, beta, ignore_index, has_label): else: n_non_ignore = BT - _jsd_kernel[(n_rows,)]( - X_ptr=_input, # input in logspace, X = log Q - X_stride=_input.stride(-2), - Y_ptr=target, # ground truth in logspace, Y = log P - Y_stride=target.stride(-2), - loss_ptr=loss, - loss_stride=loss.stride(-2), - dX_ptr=dX, - dX_stride=dX.stride(-2), - label_ptr=(shift_labels if has_label else torch.empty(1, device=_input.device)), # dummy ptr if no label - beta=beta, - n_non_ignore=n_non_ignore, - ignore_index=ignore_index, - n_cols=V, - BLOCK_SIZE=BLOCK_SIZE, - HAS_LABEL=has_label, - ) + with device_context(_input.device): + _jsd_kernel[(n_rows,)]( + X_ptr=_input, # input in logspace, X = log Q + X_stride=_input.stride(-2), + Y_ptr=target, # ground truth in logspace, Y = log P + Y_stride=target.stride(-2), + loss_ptr=loss, + loss_stride=loss.stride(-2), + dX_ptr=dX, + dX_stride=dX.stride(-2), + label_ptr=(shift_labels if has_label else torch.empty(1, device=_input.device)), # dummy ptr if no label + beta=beta, + n_non_ignore=n_non_ignore, + ignore_index=ignore_index, + n_cols=V, + BLOCK_SIZE=BLOCK_SIZE, + HAS_LABEL=has_label, + ) loss = torch.sum(loss) return loss.to(_input.dtype), dX diff --git a/liger-kernels/torch-ext/liger_kernels/kl_div.py b/liger-kernels/torch-ext/liger_kernels/kl_div.py index 706451c8..f0c4c5e0 100644 --- a/liger-kernels/torch-ext/liger_kernels/kl_div.py +++ b/liger-kernels/torch-ext/liger_kernels/kl_div.py @@ -7,6 +7,7 @@ from .utils import ensure_contiguous from .utils import is_hip from .utils import infer_device +from .utils import device_context def get_num_warps(BLOCK_SIZE): @@ -130,20 +131,21 @@ def kldiv_forward_triton(y_pred, y_true, log_target, reduction, eps): # [BT, V] out_size = (BT, V) if reduction == _REDUCTION_MODE_NONE.value else (BT,) output_tensor = torch.zeros(out_size, device=y_pred.device, dtype=torch.float32) - _kldiv_kernel_forward[grid]( - y_pred, - y_pred.stride(0), - y_true, - y_true.stride(0), - output_tensor, - output_tensor.stride(0), - V, - eps=eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - log_target=log_target, - reduction=reduction, - ) + with device_context(y_pred.device): + _kldiv_kernel_forward[grid]( + y_pred, + y_pred.stride(0), + y_true, + y_true.stride(0), + output_tensor, + output_tensor.stride(0), + V, + eps=eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + log_target=log_target, + reduction=reduction, + ) # calculated according to the reduction mode same as in Pytorch. In the later versions, `mean` will be changed to the same behavior as `batchmean` # https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html @@ -165,17 +167,18 @@ def kldiv_backward_triton(target, grad_output, new_grads, log_target): grid = (BT,) - # We store the gradients in-place in the input tensor - _kldiv_kernel_backward[grid]( - target, - target.stride(0), - new_grads, - new_grads.stride(0), - V, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - log_target=log_target, - ) + with device_context(target.device): + # We store the gradients in-place in the input tensor + _kldiv_kernel_backward[grid]( + target, + target.stride(0), + new_grads, + new_grads.stride(0), + V, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + log_target=log_target, + ) # If cross entropy is the last layer, grad_output is 1.0. Skip the mul then. if torch.equal(grad_output, torch.tensor(1.0, device=grad_output.device)): diff --git a/liger-kernels/torch-ext/liger_kernels/layer_norm.py b/liger-kernels/torch-ext/liger_kernels/layer_norm.py index 6422c3dd..90d7cbf2 100644 --- a/liger-kernels/torch-ext/liger_kernels/layer_norm.py +++ b/liger-kernels/torch-ext/liger_kernels/layer_norm.py @@ -11,6 +11,8 @@ from .utils import get_npu_core_count from .utils import set_large_grf_mode from .utils import is_npu_available +from .utils import device_context + if compare_version("triton", operator.ge, "3.0.0") and not is_npu_available(): try: @@ -202,26 +204,27 @@ def layer_norm_forward(X, W, B, eps): if X.device.type == "xpu": set_large_grf_mode(kernel_args) - # Launch kernel with one thread block per row for optimal performance - grid = (n_rows,) - _layer_norm_forward_kernel[grid]( - Y, - Y.stride(0), - X, - X.stride(0), - W, - W.stride(0), - B, - B.stride(0), - Mean, - Mean.stride(0), - RSTD, - RSTD.stride(0), - n_cols, - eps, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - **kernel_args, + with device_context(X.device): + # Launch kernel with one thread block per row for optimal performance + grid = (n_rows,) + _layer_norm_forward_kernel[grid]( + Y, + Y.stride(0), + X, + X.stride(0), + W, + W.stride(0), + B, + B.stride(0), + Mean, + Mean.stride(0), + RSTD, + RSTD.stride(0), + n_cols, + eps, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=num_warps, + **kernel_args, ) return Y.view(*shape), X, Mean, RSTD, BLOCK_SIZE, num_warps @@ -273,29 +276,30 @@ def layer_norm_backward(dY, X, W, B, Mean, RSTD): kernel_args.update({"num_warps": 32, "num_stages": 4}) set_large_grf_mode(kernel_args) - # Launch kernel with one thread block per row for optimal performance - _layer_norm_backward_kernel[grid]( - X, - X.stride(0), - W, - Mean, - Mean.stride(0), - RSTD, - RSTD.stride(0), - DX, - DX.stride(0), - _DW, - _DW.stride(0), - _DB, - _DB.stride(0), - dY, - dY.stride(0), - n_rows, - n_cols, - rows_per_program=rows_per_program, - BLOCK_SIZE=BLOCK_SIZE, - **kernel_args, - ) + with device_context(X.device): + # Launch kernel with one thread block per row for optimal performance + _layer_norm_backward_kernel[grid]( + X, + X.stride(0), + W, + Mean, + Mean.stride(0), + RSTD, + RSTD.stride(0), + DX, + DX.stride(0), + _DW, + _DW.stride(0), + _DB, + _DB.stride(0), + dY, + dY.stride(0), + n_rows, + n_cols, + rows_per_program=rows_per_program, + BLOCK_SIZE=BLOCK_SIZE, + **kernel_args, + ) DX = DX.view(*shape) DW = _DW.sum(dim=0).to(W.dtype) diff --git a/liger-kernels/torch-ext/liger_kernels/qwen2vl_mrope.py b/liger-kernels/torch-ext/liger_kernels/qwen2vl_mrope.py index fbd120f9..cdbd7255 100644 --- a/liger-kernels/torch-ext/liger_kernels/qwen2vl_mrope.py +++ b/liger-kernels/torch-ext/liger_kernels/qwen2vl_mrope.py @@ -2,6 +2,8 @@ import triton import triton.language as tl +from .utils import device_context + @triton.jit def _triton_qwen2vl_mrope( @@ -128,24 +130,25 @@ def qwen2vl_mrope_forward(q, k, cos, sin, mrope_section): cos = cos.contiguous() sin = sin.contiguous() - _triton_qwen2vl_mrope[(n_row,)]( - q, - k, - cos, - sin, - seq_len, - batch_size, - n_q_head, - n_kv_head, - head_dim, - pad_n_q_head, - pad_n_kv_head, - pad_hd, - mrope_section[0], - mrope_section[1], - BLOCK_SIZE=BLOCK_SIZE, - BACKWARD_PASS=False, - ) + with device_context(q.device): + _triton_qwen2vl_mrope[(n_row,)]( + q, + k, + cos, + sin, + seq_len, + batch_size, + n_q_head, + n_kv_head, + head_dim, + pad_n_q_head, + pad_n_kv_head, + pad_hd, + mrope_section[0], + mrope_section[1], + BLOCK_SIZE=BLOCK_SIZE, + BACKWARD_PASS=False, + ) return q.transpose(1, 2), k.transpose(1, 2), cos, sin @@ -166,25 +169,26 @@ def qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section): dq = dq.contiguous() dk = dk.contiguous() - # backward is similar to forward except swapping few ops - _triton_qwen2vl_mrope[(n_row,)]( - dq, - dk, - cos, - sin, - seq_len, - batch_size, - n_q_head, - n_kv_head, - head_dim, - pad_n_q_head, - pad_n_kv_head, - pad_hd, - mrope_section[0], - mrope_section[1], - BLOCK_SIZE=BLOCK_SIZE, - BACKWARD_PASS=True, - ) + with device_context(dq.device): + # backward is similar to forward except swapping few ops + _triton_qwen2vl_mrope[(n_row,)]( + dq, + dk, + cos, + sin, + seq_len, + batch_size, + n_q_head, + n_kv_head, + head_dim, + pad_n_q_head, + pad_n_kv_head, + pad_hd, + mrope_section[0], + mrope_section[1], + BLOCK_SIZE=BLOCK_SIZE, + BACKWARD_PASS=True, + ) return dq.transpose(1, 2), dk.transpose(1, 2) diff --git a/liger-kernels/torch-ext/liger_kernels/rope.py b/liger-kernels/torch-ext/liger_kernels/rope.py index bd8ded73..639ff81b 100644 --- a/liger-kernels/torch-ext/liger_kernels/rope.py +++ b/liger-kernels/torch-ext/liger_kernels/rope.py @@ -2,6 +2,8 @@ import triton import triton.language as tl +from .utils import device_context + @triton.jit def _triton_rope( @@ -134,27 +136,28 @@ def rope_forward(q, k, cos, sin): sin = sin.contiguous() cos_batch_size = cos.shape[0] - _triton_rope[(n_row,)]( - q, - q.stride(1), - k, - k.stride(1), - cos, - cos.stride(-2), - sin, - sin.stride(-2), - seq_len, - batch_size, - cos_batch_size, - n_q_head, - n_kv_head, - head_dim, - pad_n_q_head, - pad_n_kv_head, - pad_hd, - BLOCK_SIZE=BLOCK_SIZE, - BACKWARD_PASS=False, - ) + with device_context(q.device): + _triton_rope[(n_row,)]( + q, + q.stride(1), + k, + k.stride(1), + cos, + cos.stride(-2), + sin, + sin.stride(-2), + seq_len, + batch_size, + cos_batch_size, + n_q_head, + n_kv_head, + head_dim, + pad_n_q_head, + pad_n_kv_head, + pad_hd, + BLOCK_SIZE=BLOCK_SIZE, + BACKWARD_PASS=False, + ) return q.transpose(1, 2), k.transpose(1, 2), cos, sin @@ -176,28 +179,29 @@ def rope_backward(dq, dk, cos, sin): dq = dq.contiguous() dk = dk.contiguous() - # backward is similar to forward except swapping few ops - _triton_rope[(n_row,)]( - dq, - dq.stride(1), - dk, - dk.stride(1), - cos, - cos.stride(-2), - sin, - sin.stride(-2), - seq_len, - batch_size, - cos_batch_size, - n_q_head, - n_kv_head, - head_dim, - pad_n_q_head, - pad_n_kv_head, - pad_hd, - BLOCK_SIZE=BLOCK_SIZE, - BACKWARD_PASS=True, - ) + with device_context(dq.device): + # backward is similar to forward except swapping few ops + _triton_rope[(n_row,)]( + dq, + dq.stride(1), + dk, + dk.stride(1), + cos, + cos.stride(-2), + sin, + sin.stride(-2), + seq_len, + batch_size, + cos_batch_size, + n_q_head, + n_kv_head, + head_dim, + pad_n_q_head, + pad_n_kv_head, + pad_hd, + BLOCK_SIZE=BLOCK_SIZE, + BACKWARD_PASS=True, + ) return dq.transpose(1, 2), dk.transpose(1, 2) diff --git a/liger-kernels/torch-ext/liger_kernels/tvd.py b/liger-kernels/torch-ext/liger_kernels/tvd.py index 21b83c7c..e6ed753a 100644 --- a/liger-kernels/torch-ext/liger_kernels/tvd.py +++ b/liger-kernels/torch-ext/liger_kernels/tvd.py @@ -6,6 +6,7 @@ import triton.language as tl from .utils import ensure_contiguous +from .utils import device_context MAX_FUSED_SIZE = 65536 // 4 @@ -124,24 +125,25 @@ def tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_ else: scale = 1.0 - _tv_distance_kernel[grid]( - p, - p.stride(0), - q, - q.stride(0), - output_tensor, - output_tensor.stride(0), - grads, - grads.stride(0), - shift_labels if has_label else torch.empty(1, device=p.device), - ignore_index, - V, - scale, - BLOCK_SIZE=BLOCK_SIZE, - HAS_LABEL=has_label, - num_warps=num_warps, - reduction=reduction, - ) + with device_context(p.device): + _tv_distance_kernel[grid]( + p, + p.stride(0), + q, + q.stride(0), + output_tensor, + output_tensor.stride(0), + grads, + grads.stride(0), + shift_labels if has_label else torch.empty(1, device=p.device), + ignore_index, + V, + scale, + BLOCK_SIZE=BLOCK_SIZE, + HAS_LABEL=has_label, + num_warps=num_warps, + reduction=reduction, + ) # Loss and gradients are already scaled inside the kernel — no separate division needed if reduction in (_REDUCTION_MODE_BATCHMEAN.value, _REDUCTION_MODE_MEAN.value):