Kernels:

kernels-community
/

liger-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 16 days ago

Commit

0213eef

verified ·

1 Parent(s): 2e7a731

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

build/torch-rocm/__init__.py +8 -2
build/torch-rocm/_ops.py +8 -2
build/torch-rocm/cross_entropy.py +54 -50
build/torch-rocm/dyt.py +16 -12
build/torch-rocm/fused_linear_cross_entropy.py +72 -69
build/torch-rocm/geglu.py +22 -18
build/torch-rocm/group_norm.py +43 -39
build/torch-rocm/jsd.py +19 -17
build/torch-rocm/kl_div.py +28 -25
build/torch-rocm/layer_norm.py +47 -43
build/torch-rocm/layers.py +174 -344
build/torch-rocm/liger_kernels/__init__.py +0 -26
build/torch-rocm/metadata.json +38 -2
build/torch-rocm/metadata.json.sigstore +1 -0
build/torch-rocm/qwen2vl_mrope.py +41 -37
build/torch-rocm/rms_norm.py +97 -92
build/torch-rocm/rope.py +47 -43
build/torch-rocm/swiglu.py +23 -20
build/torch-rocm/tiled_mlp.py +91 -14
build/torch-rocm/tvd.py +20 -18
build/torch-rocm/utils.py +12 -0

build/torch-rocm/__init__.py CHANGED Viewed

@@ -1,4 +1,10 @@
 from . import layers
-from .layers import CrossEntropyOutput, LigerForCausalLMLoss
-__all__ = ["layers", "LigerForCausalLMLoss", "CrossEntropyOutput"]

 from . import layers
+from .layers import LigerForCausalLMLoss, liger_rotary_pos_emb
+__all__ = [
+    "layers",
+    # kept for BC
+    "LigerForCausalLMLoss",
+    "liger_rotary_pos_emb",
+]

build/torch-rocm/_ops.py CHANGED Viewed

@@ -4,7 +4,13 @@ def get_backend() -> str:
     """Detect the backend by inspecting torch."""
     import torch
-    if hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"
@@ -22,7 +28,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "liger_kernels"
-    unique_id = "ab435e2"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

     """Detect the backend by inspecting torch."""
     import torch
+    if hasattr(torch.backends, "tpu"):
+        # torch_tpu sets torch.backends.tpu when it is imported (via
+        # torch's device-backend autoload), regardless of whether TPU
+        # hardware is present — analogous to torch.version.cuda being
+        # set on CUDA builds without a GPU.
+        return "tpu"
+    elif hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"
 def _find_ops_name() -> str:
     kernel_name = "liger_kernels"
+    unique_id = "a5e3c9a"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-rocm/cross_entropy.py CHANGED Viewed

@@ -11,6 +11,8 @@ from .utils import element_mul_kernel
 from .utils import is_hip
 from .utils import infer_device
 from .utils import is_npu_available
 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,
-    )
     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,
-        )
     return _input

 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:
     if target.stride(-1) != 1:
         target = target.contiguous()
+    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
     # 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:
+        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

build/torch-rocm/dyt.py CHANGED Viewed

@@ -9,6 +9,8 @@ from .utils import ensure_contiguous
 from .utils import get_npu_core_count
 from .utils import infer_device
 from .utils import is_npu_available
 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,
-    )
     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)
     if HAVE_BETA:
         db = db.sum(0).to(x.dtype)
     dg = dg.sum(0).to(gamma.dtype)

 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:
     y = torch.empty_like(x)
+    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)
     db = torch.empty(NUM_SMS, N, dtype=torch.float32, device=x.device) if HAVE_BETA else None
     dx = torch.empty_like(dy)
+    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)

build/torch-rocm/fused_linear_cross_entropy.py CHANGED Viewed

@@ -7,6 +7,7 @@ from .utils import amp_custom_fwd
 from .utils import element_mul_kernel
 from .utils import is_hip
 from .utils import infer_device
 # 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,
-        )
         # 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
             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

 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
         logits_chunk = logits_chunk.contiguous()
         target_chunk = target_chunk.contiguous()
+        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:
 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)):
+        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_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
+                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

build/torch-rocm/geglu.py CHANGED Viewed

@@ -8,6 +8,8 @@ from .utils import calculate_settings
 from .utils import compare_version
 from .utils import ensure_contiguous
 from .utils import is_npu_available
 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,
-    )
     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,
-    )
     return a.view(*ori_shape), b.view(*ori_shape)

 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:
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    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)
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    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)

build/torch-rocm/group_norm.py CHANGED Viewed

@@ -8,6 +8,8 @@ from .utils import compare_version
 from .utils import ensure_contiguous
 from .utils import infer_device
 from .utils import is_npu_available
 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,
-    )
     # 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,
-    )
     # Return tensors in the original shape
     return DX.view(*shape), DW, DB

 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:
     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)
+    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
     DB = torch.zeros((num_channels), dtype=B.dtype, device=B.device)
     triton_dtype = tl.float32 if X.dtype == torch.float32 else tl.bfloat16
+    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

build/torch-rocm/jsd.py CHANGED Viewed

@@ -6,6 +6,7 @@ import triton.language as tl
 from .utils import ensure_contiguous
 from .utils import infer_device
 @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,
-    )
     loss = torch.sum(loss)
     return loss.to(_input.dtype), dX

 from .utils import ensure_contiguous
 from .utils import infer_device
+from .utils import device_context
 @triton.jit
     else:
         n_non_ignore = BT
+    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

build/torch-rocm/kl_div.py CHANGED Viewed

@@ -7,6 +7,7 @@ import triton.language as tl
 from .utils import ensure_contiguous
 from .utils import is_hip
 from .utils import infer_device
 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,
-    )
     # 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,
-    )
     # 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)):

 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):
     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)
+    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
     grid = (BT,)
+    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)):

build/torch-rocm/layer_norm.py CHANGED Viewed

@@ -11,6 +11,8 @@ from .utils import ensure_contiguous
 from .utils import get_npu_core_count
 from .utils import set_large_grf_mode
 from .utils import is_npu_available
 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,
     )
     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,
-    )
     DX = DX.view(*shape)
     DW = _DW.sum(dim=0).to(W.dtype)

 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:
     if X.device.type == "xpu":
         set_large_grf_mode(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
         kernel_args.update({"num_warps": 32, "num_stages": 4})
         set_large_grf_mode(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)

build/torch-rocm/layers.py CHANGED Viewed

@@ -1,360 +1,144 @@
 import inspect
 from dataclasses import dataclass
-from typing import Optional, Tuple
 import torch
 import torch.nn as nn
-from .cross_entropy import LigerCrossEntropyFunction
-from .dyt import LigerDyTFunction
 from .fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
 from .geglu import LigerGELUMulFunction
-from .group_norm import LigerGroupNormFunction
-from .jsd import LigerJSDFunction
-from .kl_div import LigerKLDivLossFunction
-from .layer_norm import LigerLayerNormFunction
-from .qwen2vl_mrope import LigerQwen2VLMRopeFunction
 from .rms_norm import LigerRMSNormFunction
 from .rope import LigerRopeFunction
 from .swiglu import LigerSiLUMulFunction
 from .tiled_mlp import apply_tiled_mlp
-from .tvd import LigerTVDLossFunction
 class LigerRMSNorm(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        eps: float = 1e-6,
-        offset: float = 0.0,
-        casting_mode: str = "llama",
-        init_fn: str = "ones",
-        in_place: bool = True,
-        row_mode: Optional[bool] = None,
-        elementwise_affine: bool = True,
-    ):
-        super().__init__()
-        assert init_fn in ("ones", "zeros"), f"init_fn must be 'ones' or 'zeros', got {init_fn}"
-        self.hidden_size = hidden_size
-        self.variance_epsilon = eps
-        self.offset = offset
-        self.casting_mode = casting_mode
-        self.in_place = in_place
-        self.row_mode = row_mode
-        self.elementwise_affine = elementwise_affine
-        if elementwise_affine:
-            init = torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size)
-            self.weight = nn.Parameter(init)
-        else:
-            self.register_parameter("weight", None)
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return LigerRMSNormFunction.apply(
-            hidden_states,
-            self.weight,
-            self.variance_epsilon,
-            self.offset,
-            self.casting_mode,
-            self.in_place,
-            self.row_mode,
-        )
-    def extra_repr(self) -> str:
-        return (
-            f"{self.hidden_size}, eps={self.variance_epsilon}, offset={self.offset}, "
-            f"in_place={self.in_place}, row_mode={self.row_mode}"
-        )
-class LigerLayerNorm(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        eps: float = 1e-6,
-        bias: bool = False,
-        init_fn: str = "ones",
-    ):
-        super().__init__()
-        assert init_fn in ("ones", "zeros"), f"init_fn must be 'ones' or 'zeros', got {init_fn}"
-        self.hidden_size = hidden_size
-        self.variance_epsilon = eps
-        self.weight = nn.Parameter(torch.ones(hidden_size) if init_fn == "ones" else torch.zeros(hidden_size))
-        self.bias = nn.Parameter(torch.randn(hidden_size) if bias else torch.zeros(hidden_size))
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return LigerLayerNormFunction.apply(hidden_states, self.weight, self.bias, self.variance_epsilon)
-    def extra_repr(self) -> str:
-        return f"{self.hidden_size}, eps={self.variance_epsilon}"
-class LigerGroupNorm(nn.Module):
-    def __init__(
-        self,
-        num_channels: int,
-        num_groups: int,
-        eps: float = 1e-6,
-        bias: bool = False,
-        init_fn: str = "ones",
-    ):
-        super().__init__()
-        assert init_fn in ("ones", "zeros"), f"init_fn must be 'ones' or 'zeros', got {init_fn}"
-        assert num_channels % num_groups == 0, (
-            f"num_channels ({num_channels}) must be divisible by num_groups ({num_groups})"
-        )
-        self.num_channels = num_channels
-        self.num_groups = num_groups
-        self.variance_epsilon = eps
-        self.weight = nn.Parameter(torch.ones(num_channels) if init_fn == "ones" else torch.zeros(num_channels))
-        self.bias = nn.Parameter(torch.randn(num_channels) if bias else torch.zeros(num_channels))
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        assert hidden_states.dim() >= 3, f"Input must have at least 3 dimensions, got {hidden_states.dim()}"
-        assert hidden_states.size(1) == self.num_channels, (
-            f"Input must have {self.num_channels} channels, got {hidden_states.size(1)}"
-        )
-        return LigerGroupNormFunction.apply(
             hidden_states,
             self.weight,
-            self.bias,
-            self.num_channels,
-            self.num_groups,
             self.variance_epsilon,
-        )
-    def extra_repr(self) -> str:
-        return f"num_channels={self.num_channels}, num_groups={self.num_groups}, eps={self.variance_epsilon}"
-class LigerDyT(nn.Module):
-    def __init__(self, hidden_size: int, beta: bool = True, init_alpha: float = 0.5):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.init_alpha = init_alpha
-        self.alpha = nn.Parameter(torch.ones(1) * init_alpha)
-        self.gamma = nn.Parameter(torch.ones(hidden_size))
-        self.beta = nn.Parameter(torch.zeros(hidden_size)) if beta else None
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return LigerDyTFunction.apply(x, self.alpha, self.gamma, self.beta)
-    def extra_repr(self) -> str:
-        return f"{self.hidden_size}, init_alpha={self.init_alpha}, beta={self.beta is not None}"
-class LigerCrossEntropyLoss(nn.Module):
-    def __init__(
-        self,
-        weight: Optional[torch.Tensor] = None,
-        ignore_index: int = -100,
-        lse_square_scale: float = 0.0,
-        label_smoothing: float = 0.0,
-        reduction: str = "mean",
-        softcap: Optional[float] = None,
-    ):
-        super().__init__()
-        assert 0.0 <= label_smoothing <= 1.0, f"label_smoothing must be in [0, 1], got {label_smoothing}"
-        assert reduction in ("mean", "sum", "none"), f"reduction must be 'mean', 'sum', or 'none', got {reduction}"
-        assert softcap is None or softcap > 0, f"softcap must be > 0 or None, got {softcap}"
-        self.weight = weight
-        self.ignore_index = ignore_index
-        self.lse_square_scale = lse_square_scale
-        self.label_smoothing = label_smoothing
-        self.reduction = reduction
-        self.softcap = softcap
-    def forward(self, _input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        loss, _, _, _ = LigerCrossEntropyFunction.apply(
-            _input,
-            target,
-            self.weight,
-            self.ignore_index,
-            self.lse_square_scale,
-            self.label_smoothing,
-            self.reduction,
-            self.softcap,
-            False,
-            False,
             False,
         )
-        return loss
-class LigerFusedLinearCrossEntropyLoss(nn.Module):
-    def __init__(
-        self,
-        ce_weight: Optional[torch.Tensor] = None,
-        ignore_index: int = -100,
-        lse_square_scale: float = 0.0,
-        label_smoothing: float = 0.0,
-        reduction: str = "mean",
-        softcap: Optional[float] = None,
-        accum_dtype: Optional[torch.dtype] = None,
-        use_token_scaling: bool = False,
-    ):
-        super().__init__()
-        assert 0.0 <= label_smoothing <= 1.0, f"label_smoothing must be in [0, 1], got {label_smoothing}"
-        assert reduction in ("mean", "sum", "none"), f"reduction must be 'mean', 'sum', or 'none', got {reduction}"
-        assert softcap is None or softcap > 0, f"softcap must be > 0 or None, got {softcap}"
-        self.ce_weight = ce_weight
-        self.ignore_index = ignore_index
-        self.lse_square_scale = lse_square_scale
-        self.label_smoothing = label_smoothing
-        self.reduction = reduction
-        self.softcap = softcap
-        self.accum_dtype = accum_dtype
-        self.use_token_scaling = use_token_scaling
-    def forward(
-        self,
-        lin_weight: torch.Tensor,
-        _input: torch.Tensor,
-        target: torch.Tensor,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        loss, _, _, _ = LigerFusedLinearCrossEntropyFunction.apply(
-            _input,
-            lin_weight,
-            target,
-            bias,
-            self.ce_weight,
-            self.ignore_index,
-            self.lse_square_scale,
-            self.label_smoothing,
-            self.reduction,
-            self.softcap,
-            False,
-            self.accum_dtype,
-            self.use_token_scaling,
-            False,
-            False,
-        )
-        return loss
-class LigerJSD(nn.Module):
-    def __init__(self, beta: float = 0.5, ignore_index: int = -100):
-        super().__init__()
-        self.beta = beta
-        self.ignore_index = ignore_index
-    def forward(
-        self,
-        log_q: torch.Tensor,
-        log_p: torch.Tensor,
-        shift_labels: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        return LigerJSDFunction.apply(log_q, log_p, shift_labels, self.beta, self.ignore_index)
-class LigerKLDIVLoss(nn.KLDivLoss):
-    def __init__(self, eps: float = 1e-10, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.eps = eps
-    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
-        return LigerKLDivLossFunction.apply(y_pred, y_true, self.reduction, self.log_target, self.eps)
-class LigerTVDLoss(nn.Module):
-    def __init__(self, reduction: str = "batchmean", ignore_index: int = -100):
-        super().__init__()
-        self.reduction = reduction
-        self.ignore_index = ignore_index
-    def forward(
-        self,
-        p: torch.Tensor,
-        q: torch.Tensor,
-        shift_labels: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        return LigerTVDLossFunction.apply(p, q, shift_labels, self.reduction, self.ignore_index)
 class LigerSwiGLUMLP(nn.Module):
-    """SwiGLU MLP block. ``config`` must expose ``hidden_size``, ``intermediate_size``,
-    and ``hidden_act`` (must be ``silu`` or ``swish``)."""
-    def __init__(self, config):
-        super().__init__()
-        if config.hidden_act not in ("silu", "swish"):
-            raise ValueError(f"Activation function {config.hidden_act} not supported.")
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
 class LigerGEGLUMLP(nn.Module):
-    """GEGLU MLP block. ``config`` must expose ``hidden_size`` and ``intermediate_size``.
-    Uses the tanh approximation of GELU (matches Gemma 1/1.1/2)."""
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
-class LigerTiledGEGLUMLP(nn.Module):
     gate_proj: nn.Linear
     up_proj: nn.Linear
     down_proj: nn.Linear
-    num_shards: int
-    def _mlp_forward(self, module, x):
-        """Internal MLP forward function for tiled computation."""
-        gate = module.gate_proj(x)
-        up = module.up_proj(x)
-        return module.down_proj(LigerGELUMulFunction.apply(gate, up))
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         compute_params = [p for p in self.parameters() if p.requires_grad]
         return apply_tiled_mlp(
-            fn=self._mlp_forward,
             mlp_module=self,
             x=x,
-            num_shards=self.num_shards,
             compute_params=compute_params,
         )
-class LigerTiledSwiGLUMLP(nn.Module):
     gate_proj: nn.Linear
     up_proj: nn.Linear
     down_proj: nn.Linear
-    num_shards: int
-    def _mlp_forward(self, module, x):
-        """Internal MLP forward function for tiled computation."""
-        gate = module.gate_proj(x)
-        up = module.up_proj(x)
-        return module.down_proj(LigerSiLUMulFunction.apply(gate, up))
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         compute_params = [p for p in self.parameters() if p.requires_grad]
         return apply_tiled_mlp(
-            fn=self._mlp_forward,
             mlp_module=self,
             x=x,
-            num_shards=self.num_shards,
             compute_params=compute_params,
         )
 @dataclass
 class CrossEntropyOutput:
     loss: torch.Tensor
@@ -407,23 +191,33 @@ def liger_fused_linear_cross_entropy(
     )
 def LigerForCausalLMLoss(
-    hidden_states: torch.Tensor,
-    lm_head_weight: torch.Tensor,
     labels: torch.Tensor,
-    hidden_size: int,
     num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
     shift_labels: Optional[torch.Tensor] = None,
-    final_logit_softcapping: Optional[float] = None,
-    return_token_accuracy: bool = False,
-    return_predicted_tokens: bool = False,
     **kwargs,
 ):
-    """Drop-in replacement for ``transformers.loss.ForCausalLMLoss`` that fuses the
-    final ``lm_head`` projection with the cross-entropy loss. Returns a scalar
-    ``loss`` by default; returns a :class:`CrossEntropyOutput` when
-    ``return_token_accuracy`` or ``return_predicted_tokens`` is set."""
     applicable_params = inspect.signature(liger_fused_linear_cross_entropy).parameters
     kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
@@ -435,80 +229,116 @@ def LigerForCausalLMLoss(
     shift_labels = shift_labels.view(-1).to(hidden_states.device)
     reduction = "sum" if num_items_in_batch is not None else "mean"
-    result = liger_fused_linear_cross_entropy(
         hidden_states,
         lm_head_weight,
         shift_labels,
         reduction=reduction,
         ignore_index=ignore_index,
         softcap=final_logit_softcapping,
-        return_token_accuracy=return_token_accuracy,
-        return_predicted_tokens=return_predicted_tokens,
         **kwargs,
     )
-    if isinstance(result, CrossEntropyOutput):
-        loss = result.loss
-        token_accuracy = result.token_accuracy
-        predicted_tokens = result.predicted_tokens
-    else:
-        loss = result
-        token_accuracy = None
-        predicted_tokens = None
     if reduction == "sum":
         loss = loss / num_items_in_batch
-    if return_token_accuracy or return_predicted_tokens:
-        return CrossEntropyOutput(
-            loss=loss,
-            token_accuracy=token_accuracy,
-            predicted_tokens=predicted_tokens,
-        )
     return loss
-def liger_rotary_pos_emb(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    position_ids: Optional[torch.Tensor] = None,
-    unsqueeze_dim: int = 1,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply standard rotary positional embedding to ``q`` and ``k``."""
-    return LigerRopeFunction.apply(q, k, cos, sin, position_ids, unsqueeze_dim)
-def liger_multimodal_rotary_pos_emb(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-    mrope_section,
-    unsqueeze_dim: int = 1,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """Apply Qwen2-VL multimodal rotary positional embedding (M-RoPE) to ``q`` and ``k``."""
-    return LigerQwen2VLMRopeFunction.apply(q, k, cos, sin, mrope_section, unsqueeze_dim)
 __all__ = [
     "LigerRMSNorm",
-    "LigerLayerNorm",
-    "LigerGroupNorm",
-    "LigerDyT",
-    "LigerCrossEntropyLoss",
-    "LigerFusedLinearCrossEntropyLoss",
-    "LigerJSD",
-    "LigerKLDIVLoss",
-    "LigerTVDLoss",
     "LigerSwiGLUMLP",
     "LigerGEGLUMLP",
-    "LigerTiledGEGLUMLP",
     "LigerTiledSwiGLUMLP",
-    "CrossEntropyOutput",
-    "liger_fused_linear_cross_entropy",
-    "LigerForCausalLMLoss",
     "liger_rotary_pos_emb",
-    "liger_multimodal_rotary_pos_emb",
 ]

 import inspect
 from dataclasses import dataclass
+from typing import Optional
 import torch
 import torch.nn as nn
 from .fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction
 from .geglu import LigerGELUMulFunction
 from .rms_norm import LigerRMSNormFunction
 from .rope import LigerRopeFunction
 from .swiglu import LigerSiLUMulFunction
 from .tiled_mlp import apply_tiled_mlp
 class LigerRMSNorm(nn.Module):
+    weight: nn.Parameter
+    variance_epsilon: float
+    can_torch_compile = True
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return LigerRMSNormFunction.apply(
             hidden_states,
             self.weight,
             self.variance_epsilon,
+            0,
+            "llama",
+            # Not in place: dY can be shared with other branches (e.g. a residual add).
             False,
+            None,
         )
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+class LigerLinear(nn.Module):
+    weight: nn.Parameter
+    bias: nn.Parameter | None
+    can_torch_compile = True
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        def _forward(module, x):
+            return nn.functional.linear(x, module.weight, module.bias)
+        compute_params = [p for p in self.parameters() if p.requires_grad]
+        return apply_tiled_mlp(
+            fn=_forward,
+            mlp_module=self,
+            x=input,
+            num_shards=None,
+            compute_params=compute_params,
+        )
 class LigerSwiGLUMLP(nn.Module):
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+    can_torch_compile = True
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
 class LigerGEGLUMLP(nn.Module):
+    gate_proj: nn.Linear
+    up_proj: nn.Linear
+    down_proj: nn.Linear
+    can_torch_compile = True
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj(LigerGELUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
+class LigerTiledSwiGLUMLP(nn.Module):
     gate_proj: nn.Linear
     up_proj: nn.Linear
     down_proj: nn.Linear
+    can_torch_compile = True
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _mlp_forward(module, x):
+            """Internal MLP forward function for tiled computation."""
+            gate = module.gate_proj(x)
+            up = module.up_proj(x)
+            return module.down_proj(LigerSiLUMulFunction.apply(gate, up))
         compute_params = [p for p in self.parameters() if p.requires_grad]
         return apply_tiled_mlp(
+            fn=_mlp_forward,
             mlp_module=self,
             x=x,
+            num_shards=None,
             compute_params=compute_params,
         )
+class LigerTiledGEGLUMLP(nn.Module):
     gate_proj: nn.Linear
     up_proj: nn.Linear
     down_proj: nn.Linear
+    can_torch_compile = True
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _mlp_forward(module, x):
+            """Internal MLP forward function for tiled computation."""
+            gate = module.gate_proj(x)
+            up = module.up_proj(x)
+            return module.down_proj(LigerGELUMulFunction.apply(gate, up))
         compute_params = [p for p in self.parameters() if p.requires_grad]
         return apply_tiled_mlp(
+            fn=_mlp_forward,
             mlp_module=self,
             x=x,
+            num_shards=None,
             compute_params=compute_params,
         )
+def liger_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply standard rotary positional embedding to ``q`` and ``k``."""
+    return LigerRopeFunction.apply(q, k, cos, sin, None, unsqueeze_dim)
 @dataclass
 class CrossEntropyOutput:
     loss: torch.Tensor
     )
+# NOTE: We have this intentional graph break as we encounter issues such as IMAs and Cublas errors.
+#       We know that this is an optimized kernel already so there is less ways to
+#       fuse it either way; we rely on torch compile to go through the base model to optimize.
+@torch.compiler.disable
 def LigerForCausalLMLoss(
+    logits: None,  # to match transformers signature
     labels: torch.Tensor,
+    vocab_size: int,  # to match transformers signature
     num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
     shift_labels: Optional[torch.Tensor] = None,
+    hidden_states: torch.Tensor | None = None,
+    lm_head_weight: torch.Tensor | None = None,
+    lm_head_bias: torch.Tensor | None = None,
     **kwargs,
 ):
+    # To match signature we hide these behind the kwargs but we expect a few kwargs to exist
+    hidden_size = kwargs.pop("hidden_size", None)
+    final_logit_softcapping = kwargs.pop("final_logit_softcapping", None)
+    if hidden_size is None or hidden_states is None or lm_head_weight is None:
+        raise ValueError(
+            f"`LigerForCausalLMLoss` requires the LLM's weight (found `{lm_head_weight is not None}`),"
+            f"the last hidden state (found `{hidden_states is not None}`), and the `hidden_size`"
+            f"(found `{hidden_size is not None}`). Please make sure to pass the necessary kwargs."
+        )
     applicable_params = inspect.signature(liger_fused_linear_cross_entropy).parameters
     kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
     shift_labels = shift_labels.view(-1).to(hidden_states.device)
     reduction = "sum" if num_items_in_batch is not None else "mean"
+    loss = liger_fused_linear_cross_entropy(
         hidden_states,
         lm_head_weight,
         shift_labels,
+        bias=lm_head_bias,
         reduction=reduction,
         ignore_index=ignore_index,
         softcap=final_logit_softcapping,
+        return_token_accuracy=False,
+        return_predicted_tokens=False,
         **kwargs,
     )
     if reduction == "sum":
         loss = loss / num_items_in_batch
     return loss
+# Add torch compile support for functions - in this case it's to allow this to be used
+# but the function itself will not be compiled (see the note at the function)
+liger_rotary_pos_emb.can_torch_compile = True
+LigerForCausalLMLoss.can_torch_compile = True
+class liger_rotary_pos_emb_layer(nn.Module):
+    can_torch_compile = True
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        unsqueeze_dim: int = 1,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply standard rotary positional embedding to ``q`` and ``k``."""
+        return LigerRopeFunction.apply(q, k, cos, sin, None, unsqueeze_dim)
+class LigerForCausalLMLossLayer(nn.Module):
+    can_torch_compile = True
+    # NOTE: We have this intentional graph break as we encounter issues such as IMAs and Cublas errors.
+    #       We know that this is an optimized kernel already so there is less ways to
+    #       fuse it either way; we rely on torch compile to go through the base model to optimize.
+    @torch.compiler.disable
+    def forward(
+        self,
+        logits: None,  # to match transformers signature
+        labels: torch.Tensor,
+        vocab_size: int,  # to match transformers signature
+        num_items_in_batch: Optional[int] = None,
+        ignore_index: int = -100,
+        shift_labels: Optional[torch.Tensor] = None,
+        hidden_states: torch.Tensor | None = None,
+        lm_head_weight: torch.Tensor | None = None,
+        lm_head_bias: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        # To match signature we hide these behind the kwargs but we expect a few kwargs to exist
+        hidden_size = kwargs.pop("hidden_size", None)
+        final_logit_softcapping = kwargs.pop("final_logit_softcapping", None)
+        if hidden_size is None or hidden_states is None or lm_head_weight is None:
+            raise ValueError(
+                f"`LigerForCausalLMLoss` requires the LLM's weight (found `{lm_head_weight is not None}`),"
+                f"the last hidden state (found `{hidden_states is not None}`), and the `hidden_size`"
+                f"(found `{hidden_size is not None}`). Please make sure to pass the necessary kwargs."
+            )
+        applicable_params = inspect.signature(liger_fused_linear_cross_entropy).parameters
+        kwargs = {k: v for k, v in kwargs.items() if k in applicable_params}
+        if shift_labels is None:
+            labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
+            shift_labels = labels[..., 1:].contiguous()
+        hidden_states = hidden_states.view(-1, hidden_size)
+        shift_labels = shift_labels.view(-1).to(hidden_states.device)
+        reduction = "sum" if num_items_in_batch is not None else "mean"
+        loss = liger_fused_linear_cross_entropy(
+            hidden_states,
+            lm_head_weight,
+            shift_labels,
+            bias=lm_head_bias,
+            reduction=reduction,
+            ignore_index=ignore_index,
+            softcap=final_logit_softcapping,
+            return_token_accuracy=False,
+            return_predicted_tokens=False,
+            **kwargs,
+        )
+        if reduction == "sum":
+            loss = loss / num_items_in_batch
+        return loss
 __all__ = [
     "LigerRMSNorm",
+    "LigerLinear",
     "LigerSwiGLUMLP",
     "LigerGEGLUMLP",
     "LigerTiledSwiGLUMLP",
+    "LigerTiledGEGLUMLP",
     "liger_rotary_pos_emb",
+    "liger_rotary_pos_emb_layer",
+    "LigerForCausalLMLoss",
+    "LigerForCausalLMLossLayer",
 ]

build/torch-rocm/liger_kernels/__init__.py DELETED Viewed

@@ -1,26 +0,0 @@
-import ctypes
-import importlib.util
-import sys
-from pathlib import Path
-from types import ModuleType
-def _import_from_path(file_path: Path) -> ModuleType:
-    # We cannot use the module name as-is, after adding it to `sys.modules`,
-    # it would also be used for other imports. So, we make a module name that
-    # depends on the path for it to be unique using the hex-encoded hash of
-    # the path.
-    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
-    module_name = path_hash
-    spec = importlib.util.spec_from_file_location(module_name, file_path)
-    if spec is None:
-        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
-    module = importlib.util.module_from_spec(spec)
-    if module is None:
-        raise ImportError(f"Cannot load module {module_name} from spec")
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)  # type: ignore
-    return module
-globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-rocm/metadata.json CHANGED Viewed

@@ -1,10 +1,46 @@
 {
   "name": "liger-kernels",
-  "id": "_liger_kernels_rocm_ab435e2",
-  "version": 1,
   "license": "BSD-2-Clause",
   "python-depends": [],
   "backend": {
     "type": "rocm"
   }
 }

 {
   "name": "liger-kernels",
+  "id": "_liger_kernels_rocm_a5e3c9a",
+  "version": 4,
+  "kernels-minver": "0.14.0",
   "license": "BSD-2-Clause",
   "python-depends": [],
+  "kernel-depends": [],
   "backend": {
     "type": "rocm"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "DSZMiK0xOBiMb0JdCc2K9QH4Z6lV7L8D/fTowstf/og=",
+      "_ops.py": "bCtURpaJ/MUGhq9sttUuiQOZM4pbc8c6bvnj1pFhhTs=",
+      "cross_entropy.py": "N/GPXvZTWjsu2qlAC96kuFJ4XFI0c3Qff9YtOkqaQG0=",
+      "dyt.py": "Fa7TEHZqkfNKXaxSYQt8/MlnOUFo+WIUIkz6ZA2O/b8=",
+      "fused_linear_cross_entropy.py": "l21OYH+10ryswTwtuad+h7V6Xo/8zeSLoKDyn7ZCVxo=",
+      "geglu.py": "m1Y4vqkXPBdoSzxKOTi7SiU8NW85a2JW/2/2BQbgktg=",
+      "group_norm.py": "7EcdehRRoGaias7RiU5Qil+AYQ8ZClP7mJosl1Og8r4=",
+      "jsd.py": "yWhyIa06nzvlhcwGBQNSGrTv4bessgm4vrgqQEfNUMo=",
+      "kl_div.py": "nPN5Lb2NcIGWzk+RNxgFJ69GKyfPeRLg4ycCns7KO9U=",
+      "layer_norm.py": "+t987/DupvUHRaLK8s5V9W6ewHtFZIG5WJLD5kg0RjA=",
+      "layers.py": "o90FuFiP0vXzryK3lQG7NFWyNkfWtGyqUQ89gcL7Who=",
+      "qwen2vl_mrope.py": "uxztQJ0LlMEP4Zec/yjxVJa0/uOsB5XvjUN5rM37Rf4=",
+      "rms_norm.py": "LJfKK/kutmbztb17eC3a6OGAk39PLEAHa+hzzmcCjAA=",
+      "rope.py": "IdrugShnaNJJ9apTXlMEXMjZE0IpDZI7X23gIgYS3k4=",
+      "swiglu.py": "QSjLVgdF3a4rsPdI0B0VNTmA26M2F4Y9SSv4+303WkE=",
+      "tiled_mlp.py": "LYTCxF1DoPhcoUiIg8Q28x03+5vq1w/AKclZICc4T+Y=",
+      "tvd.py": "BEL3OibZ2tfGJRyycMTSUWah/6eDdLSV2jA7ZKcxE7s=",
+      "utils.py": "MXPc1rROeMsO7dZwR1HVLG2YHFgLvoUmNjNl+iad6x0="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "commit": "a7f0afdb29a6a3372b1d47180cc0c182454c5e3b",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "a5e3c9ae2b44550c966a5fd2109b8a224448ffe4",
+      "dirty": false
+    }
   }
 }

build/torch-rocm/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrKgAwIBAgIUTYN/dKecYuFSy8iWylJxQLztS0kwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDA2MTUwNTMwWhcNMjYxMDA2MTUxNTMwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE5whRctkwOyYT2myMamsaWQWb918sLTu9f+FAxenMOk8a357MBSH8nc6lhPPDU7KLxp/sItFJRW92l5LUOmd4FqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChhNWUzYzlhZTJiNDQ1NTBjOTY2YTVmZDIxMDliOGEyMjQ0NDhmZmU0MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChhNWUzYzlhZTJiNDQ1NTBjOTY2YTVmZDIxMDliOGEyMjQ0NDhmZmU0MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChhNWUzYzlhZTJiNDQ1NTBjOTY2YTVmZDIxMDliOGEyMjQ0NDhmZmU0MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoYTVlM2M5YWUyYjQ0NTUwYzk2NmE1ZmQyMTA5YjhhMjI0NDQ4ZmZlNDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zNzQ4MzQzNzM3OC9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhEb+UyQAABAMARzBFAiBlpnWjAK1j/eAV1IoxifdKhfRT2Mmwlp03Ooy8L4u6qgIhAIvppJ7f0Nn0nGk9mNA2LTZg41wlq1TsrxkYNNRAeIOFMAoGCCqGSM49BAMDA2cAMGQCMHmzNS6irqP9TXZ5CmdzRvcDq/v40WYgtpttGwcN1ipTXjCoz/BfDU9Fzh5PRwPaVQIwUcO8s+UJWJOJOT/i6ZCmvUMFXbajPURMxwhLFniYEvO830KjFaP7lpPqIl45zjNY"}, "tlogEntries":[{"logIndex":"3111423604", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1791299130", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQC9Le9KJb7yZPMkwxn82LSWvVmgn7zbP+gNveNHzEjbUwIhAIJ81hNyxi+gyIzpTFstlTJDL4E+TjlB9EslRwpd4GgA"}, "inclusionProof":{"logIndex":"2989519342", "rootHash":"ZXKRZskZTAHxmvpU4Z+YqJgC534Dz86n1kkFy7LOVk8=", "treeSize":"2989519354", "hashes":["SLcAl9/bXfkZXL//RWNwKQidVBfTiCxNXkWAI8u2ZQQ=", "EnVECN/wbZz+gnOtscWBFAiI0Rb1lJiNy0KQonj3DmE=", "sjd5YgDCOcIQY6cN/eOufNOP1yNL981ZLbeZfku2I6U=", "mrGs7B4XcDynwbuRtph9W963YXTuA1PAK3X8yMxBSOM=", "bSw8mfeMy+OobhgMXa1U+QpCmLuo01XxVjT0hw+x7ew=", "hIhcjmKlXXWjqsIU927AphX6HWvoKqcx1cBmXctCyKo=", "MYhjnRbQvz/V0N3A/cPTBJbJ9hZW0nWOuGoVbVRzt2s=", "ymtg2bg2BTrQggg1pxCl9vaGF8BdLRBcr/sIfsN+WOc=", "0fmw4rqMbJrndCGwHEoKIxbr8PaNt4TRBnhTb0ygVtI=", "soci8BMAmEViF0iSvDe4Fr67aKooymgKvvmIeHpviec=", "uNYli6wjU0acRO+5zE2L4CKw5brwv/CVMDggUn4/ibA=", "/gAVm/mjy6BR4W9k2QfYflcuZtS4j8uS1FGmx5LRAUg=", "qLQ8+NTfJXVeQ4hnwzGZAfv08pDXZRJriV3XQZ3Lj5w=", "ewhVzfGoc5TIIeujB6CqXn5WyT2wQVElsw8KHXcxa44=", "x9rXqkSu8+Vm4yqD2XeAwn305JZ4ZykdOwcD90FjP8o=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2989519354\nZXKRZskZTAHxmvpU4Z+YqJgC534Dz86n1kkFy7LOVk8=\n\n— rekor.sigstore.dev wNI9ajBFAiB9QQ19qkO3m70dUooeo3ZXJ7SbT34FF0mfLey4DWQA7wIhAPWUVj5rm0NBQrNUUVs58lbU/XAGcryjx1hZAiwdSjZu\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJmYjRkZjc3ZTBmMDliYjcxNmJlYzYzNDM0ZjEyMmY5YTdhNTExNDhkOTA5NGYxYzI3N2EyZGM2MjEzMWVlMmUxIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJQzR2L1JLdnZnaTBHYmpQYW00ZUxsaFJyUlJiREcvVkErbXJVeFlNL2FXSkFpQUh2dE9wVG1SdHhFNlBmMnBaeE40aHZUUjdqTi9QYTRwZ0pycEo2Y3UzOXc9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KTFowRjNTVUpCWjBsVlZGbE9MMlJMWldOWmRVWlRlVGhwVjNsc1NuaFJUSHAwVXpCcmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUVRKTlZGVjNUbFJOZDFkb1kwNU5hbGw0VFVSQk1rMVVWWGhPVkUxM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUxZDJoU1kzUnJkMDk1V1ZReWJYbE5ZVzF6WVZkUlYySTVNVGh6VEZSMU9XWXJSa0VLZUdWdVRVOXJPR0V6TlRkTlFsTklPRzVqTm14b1VGQkVWVGRMVEhod0wzTkpkRVpLVWxjNU1tdzFURlZQYldRMFJuRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmhvVGxkVmVncFplbXhvV2xSS2FVNUVVVEZPVkVKcVQxUlpNbGxVVm0xYVJFbDRUVVJzYVU5SFJYbE5hbEV3VGtSb2JWcHRWVEJOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUdoT1YxVjZXWHBzYUZwVVNtbE9SRkV4Q2s1VVFtcFBWRmt5V1ZSV2JWcEVTWGhOUkd4cFQwZEZlVTFxVVRCT1JHaHRXbTFWTUUxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUdoT1YxVjZDbGw2YkdoYVZFcHBUa1JSTVU1VVFtcFBWRmt5V1ZSV2JWcEVTWGhOUkd4cFQwZEZlVTFxVVRCT1JHaHRXbTFWTUUxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlaVkZac1RUSk5OVmxYVlhsWmFsRXdUbFJWZDFsNmF6Sk9iVVV4V20xUmVVMVVRVFZaYW1ob1RXcEpNRTVFVVRRS1dtMWFiRTVFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlRucFJORTE2VVhwT2VrMHpUME01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvUldJclZYbFJRVUZDUVUxQlVucENSa0ZwUW14d2JsZHFRVXN4YWk5bFFWWXhTVzk0YVdaa1MyaG1VbFF5VFcxM2JIQXdNd3BQYjNrNFREUjFObkZuU1doQlNYWndjRW8zWmpCT2JqQnVSMnM1YlU1Qk1reFVXbWMwTVhkc2NURlVjM0o0YTFsT1RsSkJaVWxQUmsxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKalFVMUhVVU5OU0cxNlRsTTJhWEp4VURsVVdGbzFRMjFrZWxKMlkwUnhMM1kwTUZkWlozUndkSFJIZDJOT01XbHdWRmhxUTI4S2VpOUNaa1JWT1VaNmFEVlFVbmRRWVZaUlNYZFZZMDg0Y3l0VlNsZEtUMHBQVkM5cE5scERiWFpWVFVaWVltRnFVRlZTVFhoM2FFeEdibWxaUlhaUE9Bb3pNRXRxUm1GUU4yeHdVSEZKYkRRMWVtcE9XUW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg3tnb3BObAMoUcM46d2ZCs55BzE2LF3K4Xrnplje74HUCFQD/4d5EsY0f7z6gqFEA2jjeEebl9RgPMjAyNjEwMDYxNTA1MzBaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAwNjE1MDUzMFowLwYJKoZIhvcNAQkEMSIEIE7VISON7ZDB9CszdIqpT7j3X3dUthRqMXD8+kp4maVXMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQCogfg7yTQQM24PKeQ4Kg9s3Kka7wrUcAdw4+UA4mftgZkz91jcAiHhTAve2sIBpnkCMAtig9g7oSsjJHtSeCohUkKNBB9PJ+Y7UX7fN7FIdVUyhwhW53fwl+rHJU2BVAD/Ug=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"+033fg8Ju3Fr7GNDTxIvmnpRFI2QlPHCd6LcYhMe4uE="}, "signature":"MEQCIC4v/RKvvgi0GbjPam4eLlhRrRRbDG/VA+mrUxYM/aWJAiAHvtOpTmRtxE6Pf2pZxN4hvTR7jN/Pa4pgJrpJ6cu39w=="}}

build/torch-rocm/qwen2vl_mrope.py CHANGED Viewed

@@ -2,6 +2,8 @@ import torch
 import triton
 import triton.language as tl
 @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,
-    )
     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,
-    )
     return dq.transpose(1, 2), dk.transpose(1, 2)

 import triton
 import triton.language as tl
+from .utils import device_context
 @triton.jit
 def _triton_qwen2vl_mrope(
     cos = cos.contiguous()
     sin = sin.contiguous()
+    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
     dq = dq.contiguous()
     dk = dk.contiguous()
+    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)

build/torch-rocm/rms_norm.py CHANGED Viewed

@@ -24,6 +24,8 @@ from .utils import get_npu_core_count
 from .utils import set_large_grf_mode
 from .utils import torch_to_triton_dtype
 from .utils import is_npu_available
 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
-        )
     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
-        )
     dX = dX.view(*shape)
     if elementwise_affine:

 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:
     kernel_args = {}
     if X.device.type == "xpu":
         set_large_grf_mode(kernel_args)
+    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
     if X.device.type == "xpu":
         set_large_grf_mode(kernel_args)
+    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:

build/torch-rocm/rope.py CHANGED Viewed

@@ -2,6 +2,8 @@ import torch
 import triton
 import triton.language as tl
 @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,
-    )
     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,
-    )
     return dq.transpose(1, 2), dk.transpose(1, 2)

 import triton
 import triton.language as tl
+from .utils import device_context
 @triton.jit
 def _triton_rope(
     sin = sin.contiguous()
     cos_batch_size = cos.shape[0]
+    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
     dq = dq.contiguous()
     dk = dk.contiguous()
+    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)

build/torch-rocm/swiglu.py CHANGED Viewed

@@ -4,6 +4,7 @@ import triton.language as tl
 from .utils import calculate_settings
 from .utils import ensure_contiguous
 @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,
-    )
     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,
-    )
     return a.view(*ori_shape), b.view(*ori_shape)

 from .utils import calculate_settings
 from .utils import ensure_contiguous
+from .utils import device_context
 @triton.jit
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    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)
     BLOCK_SIZE, num_warps = calculate_settings(n_cols)
+    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)

build/torch-rocm/tiled_mlp.py CHANGED Viewed

@@ -5,10 +5,32 @@ from typing import List
 from typing import Optional
 import torch
 from .utils import ensure_contiguous
 class LigerTiledMLPFunction(torch.autograd.Function):
     """
     Based on DeepSpeed's TiledMLP:
@@ -40,10 +62,15 @@ class LigerTiledMLPFunction(torch.autograd.Function):
         x: torch.Tensor,
         shards: int,
         compute_params: Optional[List[torch.nn.Parameter]] = None,
     ) -> torch.Tensor:
         ctx.fn = fn
         ctx.mlp_module = mlp_module
         ctx.shards = shards
         ctx.save_for_backward(x)
         # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
@@ -61,22 +88,38 @@ class LigerTiledMLPFunction(torch.autograd.Function):
         (x,) = ctx.saved_tensors
         mlp_module = ctx.mlp_module
         shards = ctx.shards
         x_requires_grad = x.requires_grad
         x = x.detach()
         # detach() unsets x.requires_grad, so restore it
         x.requires_grad_(x_requires_grad)
-        # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
-        hidden_size = x.shape[-1]
         x_shape_orig = x.shape
         # flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
-        x = x.view(-1, hidden_size)
-        incoming_grad = grads[0].view(-1, hidden_size)
         x_grad = torch.zeros_like(x)
         x_shards = list(torch.chunk(x, chunks=shards, dim=0))
         for i, x_shard in enumerate(x_shards):
             x_shard.requires_grad_(x_requires_grad)
@@ -85,17 +128,44 @@ class LigerTiledMLPFunction(torch.autograd.Function):
             shard_step = x_shards[i].shape[0]
             shard_offset = i * x_shards[0].shape[0]
-            x_shard.grad = x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
-            incoming_grad_shard = incoming_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
-            with torch.enable_grad():
-                output = fn(mlp_module, x_shard)
-            torch.autograd.backward(output, incoming_grad_shard)
         # unflatten
         x_grad = x_grad.view(x_shape_orig)
-        return (None, None, x_grad, None, None)
 def apply_tiled_mlp(
@@ -114,9 +184,6 @@ def apply_tiled_mlp(
         x: the input tensor with shape [bs, seqlen, hidden_size] or [seqlen, hidden_size]
         num_shards: number of shards to use. If None, automatically calculated as ceil(seqlen / hidden_size)
         compute_params: list of parameters for DeepSpeed ZeRO optimization
-    Returns:
-        output tensor with the same shape as input
     """
     if num_shards is None:
         # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size]
@@ -127,10 +194,20 @@ def apply_tiled_mlp(
     # Ensure num_shards is at least 1
     num_shards = max(1, num_shards)
     return LigerTiledMLPFunction.apply(
         fn,
         mlp_module,
         x,
         num_shards,
         compute_params,
     )

 from typing import Optional
 import torch
+import torch.distributed as dist
 from .utils import ensure_contiguous
+def _autograd_input_params(
+    compute_params: Optional[List[torch.nn.Parameter]],
+) -> List[torch.nn.Parameter]:
+    """The weights to pass to the autograd Function as explicit inputs.
+    Doing so makes autograd accumulate each weight once, at the end of the backward, instead of once
+    per shard. DDP registers one hook per parameter and marks it ready on every firing, so the
+    per-shard accumulation would trip "Expected to mark a variable ready only once".
+    Returns an empty list under ZeRO-3: partitioned weights are 0-sized placeholders outside their
+    owning module's forward, so autograd would record ``(0,)`` as the expected gradient shape and the
+    recompute would fail. Those stay on the accumulate-into-``.grad`` path, where the reduction is
+    deferred to the last shard via ``ds_grad_is_ready``.
+    """
+    if not compute_params:
+        return []
+    if any(hasattr(p, "ds_id") for p in compute_params):
+        return []
+    return [p for p in compute_params if p.requires_grad]
 class LigerTiledMLPFunction(torch.autograd.Function):
     """
     Based on DeepSpeed's TiledMLP:
         x: torch.Tensor,
         shards: int,
         compute_params: Optional[List[torch.nn.Parameter]] = None,
+        *params: torch.Tensor,
     ) -> torch.Tensor:
         ctx.fn = fn
         ctx.mlp_module = mlp_module
         ctx.shards = shards
+        ctx.compute_params = compute_params
+        # `params` arrives via ensure_contiguous, which may hand us copies; resolve the originals so
+        # torch.autograd.grad differentiates the tensors the recompute actually uses.
+        ctx.grad_params = _autograd_input_params(compute_params)
         ctx.save_for_backward(x)
         # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size] (moe experts)
         (x,) = ctx.saved_tensors
         mlp_module = ctx.mlp_module
         shards = ctx.shards
+        compute_params = ctx.compute_params
+        grad_params = ctx.grad_params
+        grad_output = grads[0]
         x_requires_grad = x.requires_grad
         x = x.detach()
         # detach() unsets x.requires_grad, so restore it
         x.requires_grad_(x_requires_grad)
+        # x.shape could be [bs, seqlen, in_features] or [seqlen, in_features] (moe experts)
+        in_features = x.shape[-1]
+        out_features = grad_output.shape[-1]
         x_shape_orig = x.shape
         # flatten bs+seqlen to avoid having stride issues when narrowing into seqlen w/ bs>1
+        # NOTE: input and output feature dimensions may differ, e.g.
+        # Linear(in_features, out_features), so flatten x and grad_output with their own last dims.
+        x = x.view(-1, in_features)
+        incoming_grad = grad_output.view(-1, out_features)
         x_grad = torch.zeros_like(x)
         x_shards = list(torch.chunk(x, chunks=shards, dim=0))
+        incoming_grad_shards = list(torch.chunk(incoming_grad, chunks=shards, dim=0))
+        # ZeRO-3 partitioned parameters carry a ds_id; collect them once so the per-shard loop only
+        # flips the ready flag. Parameters on other backends have no ds_id and are left untouched.
+        ds_params = [p for p in compute_params if hasattr(p, "ds_id")] if compute_params else []
+        # Summed over shards and returned as this Function's gradients, so autograd accumulates each
+        # weight exactly once. Empty under ZeRO-3, which keeps the per-shard accumulation below.
+        param_grads = [None] * len(grad_params)
         for i, x_shard in enumerate(x_shards):
             x_shard.requires_grad_(x_requires_grad)
             shard_step = x_shards[i].shape[0]
             shard_offset = i * x_shards[0].shape[0]
+            incoming_grad_shard = incoming_grad_shards[i]
+            if grad_params:
+                with torch.enable_grad():
+                    output = fn(mlp_module, x_shard)
+                inputs = ([x_shard] if x_requires_grad else []) + grad_params
+                # torch.autograd.grad returns the gradients instead of accumulating into .grad, so no
+                # AccumulateGrad node fires here and DDP's per-parameter hook stays untouched
+                shard_grads = torch.autograd.grad(
+                    outputs=output,
+                    inputs=inputs,
+                    grad_outputs=incoming_grad_shard,
+                    allow_unused=True,
+                )
+                if x_requires_grad:
+                    x_grad.narrow(0, shard_offset, shard_step).copy_(shard_grads[0])
+                    shard_grads = shard_grads[1:]
+                for j, shard_grad in enumerate(shard_grads):
+                    if shard_grad is None:
+                        continue
+                    # the first shard's gradient is freshly allocated, so accumulate into it in place
+                    param_grads[j] = shard_grad if param_grads[j] is None else param_grads[j].add_(shard_grad)
+            else:
+                x_shard.grad = x_grad.narrow(0, shard_offset, shard_step).view_as(x_shard)
+                # Defer DeepSpeed's reduction until the last shard has accumulated into param.grad; the flag
+                # is read by ZeRO's hook during each shard's backward, so it must be set per shard.
+                for param in ds_params:
+                    param.ds_grad_is_ready = i + 1 == len(x_shards)
+                with torch.enable_grad():
+                    output = fn(mlp_module, x_shard)
+                torch.autograd.backward(output, incoming_grad_shard)
         # unflatten
         x_grad = x_grad.view(x_shape_orig)
+        return (None, None, x_grad, None, None, *param_grads)
 def apply_tiled_mlp(
         x: the input tensor with shape [bs, seqlen, hidden_size] or [seqlen, hidden_size]
         num_shards: number of shards to use. If None, automatically calculated as ceil(seqlen / hidden_size)
         compute_params: list of parameters for DeepSpeed ZeRO optimization
     """
     if num_shards is None:
         # x.shape could be [bs, seqlen, hidden_size] or [seqlen, hidden_size]
     # Ensure num_shards is at least 1
     num_shards = max(1, num_shards)
+    # All ranks must run the same number of shards: a sharded-parameter backend (DeepSpeed ZeRO-3, FSDP)
+    # gathers weights inside each shard's recompute, so a rank that runs fewer shards stops participating
+    # in those collectives and deadlocks the others. Harmonize on the per-rank maximum.
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
+        num_shards_tensor = torch.tensor(num_shards, device=x.device)
+        dist.all_reduce(num_shards_tensor, op=dist.ReduceOp.MAX)
+        num_shards = int(num_shards_tensor.item())
     return LigerTiledMLPFunction.apply(
         fn,
         mlp_module,
         x,
         num_shards,
         compute_params,
+        # passed as explicit Function inputs so autograd accumulates each weight once per backward
+        *_autograd_input_params(compute_params),
     )

build/torch-rocm/tvd.py CHANGED Viewed

@@ -6,6 +6,7 @@ import triton
 import triton.language as tl
 from .utils import ensure_contiguous
 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,
-    )
     # Loss and gradients are already scaled inside the kernel — no separate division needed
     if reduction in (_REDUCTION_MODE_BATCHMEAN.value, _REDUCTION_MODE_MEAN.value):

 import triton.language as tl
 from .utils import ensure_contiguous
+from .utils import device_context
 MAX_FUSED_SIZE = 65536 // 4
     else:
         scale = 1.0
+    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):

build/torch-rocm/utils.py CHANGED Viewed

@@ -21,6 +21,7 @@ import triton
 import triton.language as tl
 from packaging.version import Version
 def is_npu_available() -> bool:
@@ -174,3 +175,14 @@ 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"

 import triton.language as tl
 from packaging.version import Version
+from contextlib import contextmanager
 def is_npu_available() -> bool:
     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)
+    # Skipped while compiling: torch 2.14 regression, https://github.com/pytorch/pytorch/issues/199919
+    if not torch.compiler.is_compiling() and backend is not None and hasattr(backend, "device"):
+        with backend.device(device):
+            yield
+    else:
+        yield