Kernels:

kernels-community
/

liger-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 5 days ago

Commit

5c41ca4

verified ·

1 Parent(s): 7c16513

Uploaded using `kernel-builder`.

Browse files

Files changed (17) hide show

build/torch-rocm/_ops.py +1 -1
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/metadata.json +17 -17
build/torch-rocm/metadata.json.sigstore +1 -1
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/tvd.py +20 -18
build/torch-rocm/utils.py +11 -0

build/torch-rocm/_ops.py CHANGED Viewed

@@ -22,7 +22,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "liger_kernels"
-    unique_id = "4d9f798"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "liger_kernels"
+    unique_id = "0c5fb33"
     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/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "liger-kernels",
-  "id": "_liger_kernels_rocm_4d9f798",
   "version": 3,
   "license": "BSD-2-Clause",
   "python-depends": [],
@@ -11,24 +11,24 @@
     "algorithm": "sha256",
     "files": {
       "__init__.py": "DSZMiK0xOBiMb0JdCc2K9QH4Z6lV7L8D/fTowstf/og=",
-      "_ops.py": "pdeXrc8M9KCvCpe0mtloEEFT3iJWZab9CMNsgT2CK10=",
-      "cross_entropy.py": "LsQN0OI8VHqgUObWF0qqTgNA9z7SYMaZNb2I7FIYnO4=",
-      "dyt.py": "806Tai0T+Yi7Ctf5YGmtwAc87MUjlzjsLky7b1T+oWM=",
-      "fused_linear_cross_entropy.py": "yjtZHWyBcXCBJ57NbbrQ1wMmTr9nARqDBrdjXVYpR74=",
-      "geglu.py": "Lmw0I38IBIHkNROTf9V3bP6PZLktCWCqNAxzVRaipbk=",
-      "group_norm.py": "xdY/tCr4rHOb4PhD827259XB7QDb22z2psUKzTYTKBE=",
-      "jsd.py": "CftyIN36AU63e1GdtuA7hNhs66I3TjS9lttcFriBnGA=",
-      "kl_div.py": "VQG/89mefLifPeVEv+1O8hmsXgGujMLA0JXDwO5QvtU=",
-      "layer_norm.py": "INpxK7Ac27f+eckcMfOOSgZErzDhNF5iGt207eV9aoM=",
       "layers.py": "+50F9xmnpQXbU1K/lZxdJwbnTotkksJlhoLOWHjo+/M=",
       "liger_kernels/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
-      "qwen2vl_mrope.py": "3GExhYpLgB4VUtyZyjRk8XjEur3W4EWF6HQ67ML5vBU=",
-      "rms_norm.py": "GLzIZjDP74QxaeKQ5B/VDXKTGxD6Xo9cVaW/iqmsc4c=",
-      "rope.py": "v+7JHRrv+5ImoROkpKfl30WwWI4qTa2tAl7zQeB4ml4=",
-      "swiglu.py": "TLhyvpr7nbVbMqWbx8uz+iU2GxS6q0QrVxWAd1eQCBo=",
       "tiled_mlp.py": "PJw2R8YdHHxQRSwbFJWjcHxgzssEqsRWqyogkY8pKR8=",
-      "tvd.py": "VbWUtpJ79E2atqjBoVyTSO785aS+idHiQBIrJwifDAw=",
-      "utils.py": "Y8wQsplz6ed7oPmBfX7yGGcX74VB3t/e5VFvYv/X7wU="
     }
   },
   "provenance": {
@@ -38,7 +38,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "4d9f798b95ed280b770213048b8d3f41509d0c9e",
       "dirty": false
     }
   }

 {
   "name": "liger-kernels",
+  "id": "_liger_kernels_rocm_0c5fb33",
   "version": 3,
   "license": "BSD-2-Clause",
   "python-depends": [],
     "algorithm": "sha256",
     "files": {
       "__init__.py": "DSZMiK0xOBiMb0JdCc2K9QH4Z6lV7L8D/fTowstf/og=",
+      "_ops.py": "Z4DkRS77DKP3IQGbv+oQOt0NNFzokK2hpXfcHoNpG/I=",
+      "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": "+50F9xmnpQXbU1K/lZxdJwbnTotkksJlhoLOWHjo+/M=",
       "liger_kernels/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "qwen2vl_mrope.py": "uxztQJ0LlMEP4Zec/yjxVJa0/uOsB5XvjUN5rM37Rf4=",
+      "rms_norm.py": "LJfKK/kutmbztb17eC3a6OGAk39PLEAHa+hzzmcCjAA=",
+      "rope.py": "IdrugShnaNJJ9apTXlMEXMjZE0IpDZI7X23gIgYS3k4=",
+      "swiglu.py": "QSjLVgdF3a4rsPdI0B0VNTmA26M2F4Y9SSv4+303WkE=",
       "tiled_mlp.py": "PJw2R8YdHHxQRSwbFJWjcHxgzssEqsRWqyogkY8pKR8=",
+      "tvd.py": "BEL3OibZ2tfGJRyycMTSUWah/6eDdLSV2jA7ZKcxE7s=",
+      "utils.py": "l4/OveROs0DFMcRaJsaoiJBKq7jvHbAL9aRRrkZb+iM="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "sha": "0c5fb33d21084e07750a248da12071ff011bd70a",
       "dirty": false
     }
   }

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

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtGgAwIBAgIUeKKjysjZHPTkcjmFb3CHS1TysN0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzE0MTIwMjUwWhcNMjYwNzE0MTIxMjUwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEMw0GrkZBBP9oqnpDe5wf8q41n6WefaLaklEXON/6Xd4W4HwqUuA5A2sreVnQq1EjKZd4nZvn0yIbO+G/5kGEfKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUtZx+9iAL0iQdwe1N8MZBeaC2LFIwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoNGQ5Zjc5OGI5NWVkMjgwYjc3MDIxMzA0OGI4ZDNmNDE1MDlkMGM5ZTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoNGQ5Zjc5OGI5NWVkMjgwYjc3MDIxMzA0OGI4ZDNmNDE1MDlkMGM5ZTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoNGQ5Zjc5OGI5NWVkMjgwYjc3MDIxMzA0OGI4ZDNmNDE1MDlkMGM5ZTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDRkOWY3OThiOTVlZDI4MGI3NzAyMTMwNDhiOGQzZjQxNTA5ZDBjOWUwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjkzMzA2MjM4NjUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn2CCKbEAAAQDAEcwRQIgMoVU3dtBnGlxMMb3sOTM8i/rOdiqaQ37x0tz9NOZa+MCIQCvrHRw6fB5k7rokMe6iDv+n9siBeeEDI5Helc4t8PNfzAKBggqhkjOPQQDAwNpADBmAjEAuazcBiTcv8K6MyCbxJkO8Q1WX752+NAYF6uQhREhOlfLY1Htf2AgU5YzxWwv72nSAjEAsyIagTrwrpp0hqNOQI9ik7Cl4rMqprpQ90wKXEz5+urqVmetTmkd7YsMArDzriDo"}, "tlogEntries":[{"logIndex":"2167855782", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1784030571", "inclusionPromise":{"signedEntryTimestamp":"MEQCIFPwhR+59GwljQ+ZxD1rhaNTV3Sh4vPwQeU/9JrQoO2xAiBPUFa1Dck8XvuXhEpd/t74KU9P4qaGczgDzl//4hAk8w=="}, "inclusionProof":{"logIndex":"2045951520", "rootHash":"Re/Ph87BOM7jhvV4VE+3WsyQb72tumWdlqTxAOe0qEM=", "treeSize":"2045951527", "hashes":["nUla+MAuN3xIPh3Fw90o+gJvejDEBaNrSVpncG6U3kM=", "dEagCZIabgWBH6ViR5uwn8xXp3WdvaR3vEy+Q8US9b8=", "4tMZ/2/4pBxGx11gMAczICVOHyul021sgz6LWWWRpIQ=", "m+Vk8U0/4EcekuSwkXtYOJcYVvuBCDXI7XChKnMHSec=", "iMBom1ahASM6ZAxymZoy+ISJTz5W3t4BNlJ+0yiYrnc=", "IuHNjKNGzZgw8Ec1vaQ4yxMC8m5mfRHL/sO7dTeEGIM=", "T/vXn6bo+SCijMlRZwcNRxHpWcvInutRlkDvk1Y7oEM=", "TXWJIUMj+J24/w5Nh86FuokZClxMAnvB6sia+ArBfNI=", "pfPn/gRvLHfxa/LNybYhqG78qHZ4QD3PI+zQ35oYweM=", "sMaCrs8erYNQ48jT2FC56F0RAbpTkuSXZ+cs1zcSObw=", "zGam2EiRYwS9IjemppqM/acVfU15hw3eAqn9Fb8+8vY=", "qQGtNWYnFxNQ1OmqNBR2y92EFrgo5lyaKNQLaE/Dl1c=", "/UAkaaq9D8BPMw5qLAxaf+ur4Sfbv3FfLWB2M/py3wM=", "56+1wf4OgrBZEzfW9n13hPfm0bRvMRgerPv4TPE+3Ho=", "mTSsSZdzVhRgljy9CvDq3GfjxxeOTHrwgpOWK4rRj3Q=", "5sxNZDoxEj6DMmQATisX9bQXdFmeRHYfC8BjyDIIgec=", "Rq/A2aTC9e54ldNjcpsJ26rX/h8JlN6ZHxgw7yEa6C0=", "+/VZ56MsIPxMiyLAodzKXo5TEWdQp36z89qLhpzloAo=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2045951527\nRe/Ph87BOM7jhvV4VE+3WsyQb72tumWdlqTxAOe0qEM=\n\n— rekor.sigstore.dev wNI9ajBEAiAxviZPcApj0ZJ1BUlQmvWCHwh/I05iSQg0PsjajUUu5wIgfTmpusfaD0Tu2YE2XuUcR3X1uFqLW7goWine86qQDnQ=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiNmJkMTIwYTVhOGNkYWE0MmI5MTkxYWYwYzBlZDQzNDMxMDE5OWU2Y2RkZDg1YzFhNTE4NTAxODE3Yzc2NDliIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNHMmE1Y2pEL09zY3UrdlpxWjVnbXB6S2xVTGhXZHhyRmVSSE00dlRJa0p3SWhBSyttWkE1WlBZUGwrZnJUSFdLYk5TZUxFaFh5TjdzZ3VCNmpMMWxBTEhHTCIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5SSFowRjNTVUpCWjBsVlpVdExhbmx6YWxwSVVGUnJZMnB0Um1JelEwaFRNVlI1YzA0d2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2UlRCTlZFbDNUV3BWZDFkb1kwNU5hbGwzVG5wRk1FMVVTWGhOYWxWM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZOZHpCSGNtdGFRa0pRT1c5eGJuQkVaVFYzWmpoeE5ERnVObGRsWm1GTVlXdHNSVmdLVDA0dk5saGtORmMwU0hkeFZYVkJOVUV5YzNKbFZtNVJjVEZGYWt0YVpEUnVXblp1TUhsSllrOHJSeTgxYTBkRlprdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlYwV25nckNqbHBRVXd3YVZGa2QyVXhUamhOV2tKbFlVTXlURVpKZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU9SMUUxV21wak5VOUhTVFZPVjFaclRXcG5kMWxxWXpOTlJFbDRDazE2UVRCUFIwazBXa1JPYlU1RVJURk5SR3hyVFVkTk5WcFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDVIVVRWYWFtTTFUMGRKTlU1WFZtdE5hbWQzV1dwak0wMUVTWGhOZWtFd1QwZEpORnBFVG0wS1RrUkZNVTFFYkd0TlIwMDFXbFJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDVIVVRWYWFtTTFUMGRKTlU1WFZtdE5hbWQzV1dwak0wMUVTWGdLVFhwQk1FOUhTVFJhUkU1dFRrUkZNVTFFYkd0TlIwMDFXbFJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUlNhMDlYV1RNS1QxUm9hVTlVVm14YVJFazBUVWRKTTA1NlFYbE5WRTEzVGtSb2FVOUhVWHBhYWxGNFRsUkJOVnBFUW1wUFYxVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW10NlRYcEJNazFxVFRST2FsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNHlRME5MWWtWQlFVRlJSQXBCUldOM1VsRkpaMDF2VmxVelpIUkNia2RzZUUxTllqTnpUMVJOT0drdmNrOWthWEZoVVRNM2VEQjBlamxPVDFwaEswMURTVkZEZG5KSVVuYzJaa0kxQ21zM2NtOXJUV1UyYVVSMksyNDVjMmxDWldWRlJFazFTR1ZzWXpSME9GQk9abnBCUzBKblozRm9hMnBQVUZGUlJFRjNUbkJCUkVKdFFXcEZRWFZoZW1NS1FtbFVZM1k0U3paTmVVTmllRXByVHpoUk1WZFlOelV5SzA1QldVWTJkVkZvVWtWb1QyeG1URmt4U0hSbU1rRm5WVFZaZW5oWGQzWTNNbTVUUVdwRlFRcHplVWxoWjFSeWQzSndjREJvY1U1UFVVazVhV3MzUTJ3MGNrMXhjSEp3VVRrd2QwdFlSWG8xSzNWeWNWWnRaWFJVYld0a04xbHpUVUZ5UkhweWFVUnZDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg/3UMGV8fwrsfiQEAzqbUSlYc87EUEy8wrOfhaOM5NwACFBwa0r74wyBkjdrXOXL45py8exMrGA8yMDI2MDcxNDEyMDI1MVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdwwggHYAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNzE0MTIwMjUxWjAvBgkqhkiG9w0BCQQxIgQgC5Qp2/4m1NXACJdcq25HSjgek+k/yuqiVNfBceqic8kwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGgwZgIxALBGy+9QLpTAvIjqpcf3WnWsVHfTrPKGAClJboW2OzfnuQ9wBO4Qi9+XIsmEIuGPdwIxAKbZ0fUnZ4D5AG3RMAvGaeLMf2WAwCkkT+7PEpjMl7uXe2luLTuwr+SqoAxPJxE2KA=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"tr0SClqM2qQrkZGvDA7UNDEBmebN3YXBpRhQGBfHZJs="}, "signature":"MEYCIQCG2a5cjD/Oscu+vZqZ5gmpzKlULhWdxrFeRHM4vTIkJwIhAK+mZA5ZPYPl+frTHWKbNSeLEhXyN7sguB6jL1lALHGL"}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUO6RSJ/EnqV44CIEDdOOD20lhcqIwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNzIwMTMzMzA2WhcNMjYwNzIwMTM0MzA2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE8cQGo4I4BMQ8MJoUyiqSGUWjUiCqFqm2tJO888SUpnLtnkFo1FzBQjCHpomiYGRFWwqO0B19Wr1fIelQoxsFUKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUv6O4N4tqJqgSLVrYxeRJ/tmXx2UwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMGM1ZmIzM2QyMTA4NGUwNzc1MGEyNDhkYTEyMDcxZmYwMTFiZDcwYTATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMGM1ZmIzM2QyMTA4NGUwNzc1MGEyNDhkYTEyMDcxZmYwMTFiZDcwYTAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMGM1ZmIzM2QyMTA4NGUwNzc1MGEyNDhkYTEyMDcxZmYwMTFiZDcwYTAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDBjNWZiMzNkMjEwODRlMDc3NTBhMjQ4ZGExMjA3MWZmMDExYmQ3MGEwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjk3NDYyMzc1MzUvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn3+68ykAAAQDAEcwRQIhAKfi0uLcjC+7sKlx6lksgEPD/NhuLlgiGj7ibyikciwIAiBL8jxz6LA5Ggrj5nN1SaUm/ttST2HZx2F8xHAd1hu32zAKBggqhkjOPQQDAwNoADBlAjBowDB5OdvLqQ2GqWrG8F+E7lNV2bbvOsz5keEUywdb1JKbKczCW+ZE5TPgasmM0skCMQC5VBemdmpIL4GaKEgobQkePdAc/wQcmTF743v/+wJ6rtHl1qO6tRVmurEflsSKryQ="}, "tlogEntries":[{"logIndex":"2206355914", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1784554386", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDG3iFXAb+fjoVxsKm13pM8c5eEoZl7sIUSR7PCOJQPJwIhAMNN6Gc2ZifBnxZxAFkqOD6fZ9pxz6eIigLvWieBe+12"}, "inclusionProof":{"logIndex":"2084451652", "rootHash":"zbQwyp11BrPe3wggHlyhiEtVAaHleIQPjvxj/sBE5zI=", "treeSize":"2084451654", "hashes":["499kcb89L18pO0S0b4IQogj4NQ5pOzw5FHU004db5jY=", "sd8iJSY7iV+Ikc42LJKo8JVRb2fypg73SOufXNOi0Xo=", "p+vUmE/YgN+ll593XYjUalho2bs1PslkYfN9ocL4t08=", "uandAueMLs9RC3xIj77ePNadPcbKcEeORWqtfLFvAu8=", "SjVrRtIeid4t+H6u4drFM6V4zL5B3xFoI1cbwfFF6kI=", "oz8/zcwvzjAiODVBCpcZm4tSU+3KvdHqIQqxFRSS418=", "LH3MOzxultOUBpT2DD8FsK1ItkTk2kLtqOkxxthdBAE=", "NOqu2omVaTD0z/phbf6aHupVZtPpniJjH7Mohp9dQVg=", "UaVIhKZL9YPl5aIrrdhBcDNZES0kGKGGU+n9C0aUu04=", "S6fX78YW7K/Nsie8dVZY6CYd6Qnvj92xsSI5UOFLMIs=", "phX4Nc7EXYzCDd9JqHuZvmMmWgixvVF/zVGv5sTGpx4=", "V6wr2bDyz/qjnYx99CtUrvwLeHOpXEswquva8IIXU8U=", "GxT6+CEXxY8Ak/besnRW/IK3DhneHQ5X5rIcw8L01DQ=", "Rq/A2aTC9e54ldNjcpsJ26rX/h8JlN6ZHxgw7yEa6C0=", "+/VZ56MsIPxMiyLAodzKXo5TEWdQp36z89qLhpzloAo=", "daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=", "DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2084451654\nzbQwyp11BrPe3wggHlyhiEtVAaHleIQPjvxj/sBE5zI=\n\n— rekor.sigstore.dev wNI9ajBFAiAN5Kb1loICwLlR0NDXi67Zk+v9N/w+2dwjAk1gfEA7QwIhAKxiGa8YCzAVUYhlDnLZ/mDkJsC0CNlX2XHsegUBfTIa\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2Y2I2NDk4NTk2YzMwNDFiMWE5MTQzMDI5MzkwNTBmNmUxOGJiYzMzNmFjNDM4YjI3NTA3MGMyYzg1YjJkNzhlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUUNjZ3NYdHNJWWxWRTJ4YVZtcDJkdzJpdjNPZzlDK0MvbEV2NUtyZUxEUDhRSWhBTGxEWUxmaGkrY0FDL2lwUnUxaHMralA2aHdIWU5aR1FTclA5a1lFcHMyVyIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlR6WlNVMG92Ulc1eFZqUTBRMGxGUkdSUFQwUXlNR3hvWTNGSmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDU2U1hkTlZFMTZUWHBCTWxkb1kwNU5hbGwzVG5wSmQwMVVUVEJOZWtFeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVU0WTFGSGJ6UkpORUpOVVRoTlNtOVZlV2x4VTBkVlYycFZhVU54Um5GdE1uUktUemdLT0RoVFZYQnVUSFJ1YTBadk1VWjZRbEZxUTBod2IyMXBXVWRTUmxkM2NVOHdRakU1VjNJeFprbGxiRkZ2ZUhOR1ZVdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlYyTms4MENrNDBkSEZLY1dkVFRGWnlXWGhsVWtvdmRHMVllREpWZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5SMDB4V20xSmVrMHlVWGxOVkVFMFRrZFZkMDU2WXpGTlIwVjVDazVFYUd0WlZFVjVUVVJqZUZwdFdYZE5WRVpwV2tSamQxbFVRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFIVFRGYWJVbDZUVEpSZVUxVVFUUk9SMVYzVG5wak1VMUhSWGxPUkdocldWUkZlVTFFWTNnS1dtMVpkMDFVUm1sYVJHTjNXVlJCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFIVFRGYWJVbDZUVEpSZVUxVVFUUk9SMVYzVG5wak1VMUhSWGtLVGtSb2ExbFVSWGxOUkdONFdtMVpkMDFVUm1sYVJHTjNXVlJCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkNhazVYV21rS1RYcE9hMDFxUlhkUFJGSnNUVVJqTTA1VVFtaE5hbEUwV2tkRmVFMXFRVE5OVjFwdFRVUkZlRmx0VVROTlIwVjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1zelRrUlplVTE2WXpGTmVsVjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNHpLelk0ZVd0QlFVRlJSQXBCUldOM1VsRkphRUZMWm1rd2RVeGpha01yTjNOTGJIZzJiR3R6WjBWUVJDOU9hSFZNYkdkcFIybzNhV0o1YVd0amFYZEpRV2xDVERocWVIbzJURUUxQ2tkbmNtbzFiazR4VTJGVmJTOTBkRk5VTWtoYWVESkdPSGhJUVdReGFIVXpNbnBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcENiM2RFUWpVS1QyUjJUSEZSTWtkeFYzSkhPRVlyUlRkc1RsWXlZbUoyVDNONk5XdGxSVlY1ZDJSaU1VcExZa3RqZWtOWEsxcEZOVlJRWjJGemJVMHdjMnREVFZGRE5RcFdRbVZ0Wkcxd1NVdzBSMkZMUldkdllsRnJaVkJrUVdNdmQxRmpiVlJHTnpRemRpOHJkMG8yY25SSWJERnhUelowVWxadGRYSkZabXh6VTB0eWVWRTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgvw1pUaMIhxjbKLh6+5XKQorzGNdAiR35CAi3ahtL/wICFQDUaY9j5Oats6jvX/DfaOYqFh0byxgPMjAyNjA3MjAxMzMzMDZaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDcyMDEzMzMwNlowLwYJKoZIhvcNAQkEMSIEIPCocok9oFF41l0u0wS8It++csm/ItVfTR/zQ6TitjHTMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDUU5At837LJul65e/JIa4/4I1tVXnMU7HV2Y1f3HuM4ddjXQZSbWn8if7ZaoDFenoCMQCEPz2BqJsUGrxCWijUtVD8SJx79reROK3HNxmShRIsaA9ahJiE6V2W137HRpPWsR8="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"bLZJhZbDBBsakUMCk5BQ9uGLvDNqxDiydQcMLIWy144="}, "signature":"MEYCIQCcgsXtsIYlVE2xaVmp2dw2iv3Og9C+C/lEv5KreLDP8QIhALlDYLfhi+cAC/ipRu1hs+jP6hwHYNZGQSrP9kYEps2W"}}

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/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,13 @@ def set_large_grf_mode(kernel_args: dict):
     else:
         # API was changed in https://github.com/intel/intel-xpu-backend-for-triton/pull/5430
         kernel_args["grf_mode"] = "large"

 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)
+    if backend is not None and hasattr(backend, "device"):
+        with backend.device(device):
+            yield
+    else:
+        yield