quack-kernels 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

quack/cross_entropy.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
 
 import math
+from functools import partial
 from typing import Optional, Type, Literal
 
 import torch
@@ -10,10 +11,12 @@ import cuda.bindings.driver as cuda
 
 import cutlass
 import cutlass.cute as cute
-from cutlass import Int32, Float32, Boolean, const_expr
-from cutlass.cute.runtime import from_dlpack
+from cutlass import Int32, Int64, Float32, Boolean, const_expr
 
 import quack.utils as utils
+import quack.copy_utils as copy_utils
+import quack.layout_utils as layout_utils
+from quack.compile_utils import make_fake_tensor as fake_tensor
 from quack.reduce import row_reduce, online_softmax_reduce
 from quack.reduction_base import ReductionBase
 from quack.cute_dsl_utils import torch2cute_dtype_map
@@ -26,46 +29,29 @@ class CrossEntropy(ReductionBase):
             dtype,
             N,
             stage=2 if not online_softmax else 1,
-            reduction_dtype=Float32 if not online_softmax else cutlass.Int64,
+            reduction_dtype=Float32 if not online_softmax else Int64,
         )
         self.online_softmax = online_softmax
         self.reload_from = None if N <= 16384 or online_softmax else "smem"
 
-    def _calculate_threads_per_row(self):
+    def _threads_per_row(self):
         N = self.N
-        return (
-            8
-            if N <= 64
-            else (
-                16
-                if N <= 128
-                else (32 if N <= 3072 else (64 if N <= 6144 else (128 if N <= 16384 else 256)))
-            )
-        )
+        for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]:
+            if N <= limit:
+                return threads
+        return 256
 
     def _set_cluster_n(self):
         N = self.N
         if const_expr(self.dtype.width == 16):
-            cluster_n = (
-                1
-                if N <= 16 * 1024
-                else (
-                    2
-                    if N <= 32 * 1024
-                    else (4 if N <= 64 * 1024 else (8 if N <= 128 * 1024 else 16))
-                )
-            )
-        else:  # fp32
-            cluster_n = (
-                1
-                if N <= 16 * 1024
-                else (
-                    2
-                    if N <= 64 * 1024
-                    else (4 if N <= 128 * 1024 else (8 if N <= 256 * 1024 else 16))
-                )
-            )
-        self.cluster_n = cluster_n
+            thresholds = [(16 * 1024, 1), (32 * 1024, 2), (64 * 1024, 4), (128 * 1024, 8)]
+        else:
+            thresholds = [(16 * 1024, 1), (64 * 1024, 2), (128 * 1024, 4), (256 * 1024, 8)]
+        for limit, cluster in thresholds:
+            if N <= limit:
+                self.cluster_n = cluster
+                return
+        self.cluster_n = 16
 
     @cute.jit
     def __call__(
@@ -82,19 +68,30 @@ class CrossEntropy(ReductionBase):
         assert mX.element_type == self.dtype
         if const_expr(mTargetLogit is None):
             mTargetLogit = mX
+        if const_expr(mdX is not None):
+            assert mdX.element_type == self.dtype
         self._set_cluster_n()
-        # e.g. if self.N isn't divisible by 8 for bf16, we might use 64 bits (4 elements) copy
-        num_copy_bits = math.gcd(self.N, 128 // self.dtype.width) * self.dtype.width
-        tiler_mn, tv_layout = self._get_tv_layout(num_copy_bits=num_copy_bits)
-        num_threads = cute.size(tv_layout, mode=[0])
-        num_warps = num_threads // cute.arch.WARP_SIZE
+        largest_dtype_width = const_expr(mX.element_type.width)
+        if const_expr(mdX is not None):
+            largest_dtype_width = const_expr(max(largest_dtype_width, mdX.element_type.width))
+        vecsize = math.gcd(self.N, 128 // largest_dtype_width)
+        tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize)
+        num_threads = tiled_copy.size
         self.kernel(
-            mX, mTarget, mTargetLogit, mLoss, mLSE, mdX, ignore_index, tv_layout, tiler_mn
+            mX,
+            mTarget,
+            mTargetLogit,
+            mLoss,
+            mLSE,
+            mdX,
+            ignore_index,
+            tiler_mn,
+            tiled_copy,
+            threads_per_row,
         ).launch(
             grid=[cute.ceil_div(mX.shape[0], tiler_mn[0]), self.cluster_n, 1],
             block=[num_threads, 1, 1],
-            cluster=([1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None),
-            smem=self._smem_size_in_bytes(tiler_mn, num_warps),
+            cluster=[1, self.cluster_n, 1] if const_expr(self.cluster_n > 1) else None,
             stream=stream,
         )
 
@@ -108,47 +105,40 @@ class CrossEntropy(ReductionBase):
         mLSE: Optional[cute.Tensor],  # (M,)
         mdX: Optional[cute.Tensor],  # (M, N) - if provided, compute gradient
         ignore_index: Int32,  # Index to ignore in loss computation
-        tv_layout: cute.Layout,
         tiler_mn: cute.Shape,
+        tiled_copy: cute.TiledCopy,
+        threads_per_row: cutlass.Constexpr[int],
     ):
         tidx, _, _ = cute.arch.thread_idx()
         bidx, _, _ = cute.arch.block_idx()
-        if const_expr(self.cluster_n > 1):
-            cluster_y = cute.arch.block_idx()[1]
-        else:
-            cluster_y = const_expr(0)
+        cluster_y = const_expr(0) if const_expr(self.cluster_n == 1) else cute.arch.block_idx()[1]
+        tv_layout = tiled_copy.layout_tv_tiled
 
-        shape: cute.Shape = mX.shape
+        shape = mX.shape
         idX = cute.make_identity_tensor(shape)
         # slice for CTAs
-        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
-        mX_off = utils.domain_offset_i64((bidx * tiler_mn[0], 0), mX)
-        gX = cute.local_tile(mX_off, tiler_mn, (0, cluster_y))
-        cX = cute.local_tile(idX, tiler_mn, (bidx, cluster_y))
+        gX, cX = [cute.local_tile(mT, tiler_mn, (bidx, cluster_y)) for mT in (mX, idX)]
 
         smem = cutlass.utils.SmemAllocator()
         sX = smem.allocate_tensor(
-            mX.element_type,
-            cute.make_ordered_layout(tiler_mn, order=(1, 0)),
-            byte_alignment=16,
+            mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16
         )
         reduction_buffer, mbar_ptr = self._allocate_reduction_buffer_and_mbar(smem, tv_layout)
 
-        # declare the atoms which will be used later for memory copy
-        num_copy_elems_X = tv_layout.shape[1][0]
-        num_copy_bits_X = mX.element_type.width * num_copy_elems_X
-        copy_atom_load_X = cute.make_copy_atom(
-            cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=num_copy_bits_X
-        )
-        thr_copy_X = cute.make_tiled_copy(copy_atom_load_X, tv_layout, tiler_mn).get_slice(tidx)
+        thr_copy = tiled_copy.get_slice(tidx)
 
-        #### Partition to get thread view
-        tXgX = thr_copy_X.partition_S(gX)
-        tXsX = thr_copy_X.partition_D(sX)
-        tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
+        tXgX = thr_copy.partition_S(gX)
+        tXsX = thr_copy.partition_D(sX)
+        tXcX = thr_copy.partition_S(cX)[(0, None), None, None]
         tXrX = cute.make_fragment_like(tXgX)
 
-        num_warps = cute.size(tv_layout, mode=[0]) // cute.arch.WARP_SIZE
+        is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n)
+        tXpX = (
+            None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1])
+        )
+        copy = partial(copy_utils.copy, pred=tXpX)
+
+        num_warps = cute.size(tiled_copy) // cute.arch.WARP_SIZE
         self._initialize_cluster(tidx, mbar_ptr, num_warps)
 
         row = tXcX[0][0]
@@ -156,14 +146,8 @@ class CrossEntropy(ReductionBase):
         if row < shape[0]:
             target = Int32(mTarget[row])
 
-        is_even_N = const_expr(shape[1] == tiler_mn[1] * self.cluster_n)
-        tXpX = (
-            utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1])
-            if const_expr(not is_even_N)
-            else None
-        )
         if row < shape[0]:
-            cute.copy(copy_atom_load_X, tXgX, tXsX, pred=tXpX)
+            copy(tXgX, tXsX, is_async=True)
         cute.arch.cp_async_commit_group()
         cute.arch.cp_async_wait_group(0)
         # Fill OOB values with -inf
@@ -177,14 +161,11 @@ class CrossEntropy(ReductionBase):
         if row < shape[0] and tXcX[0][1] == 0 and not should_ignore:
             # Only load target logit if not ignoring this index
             if const_expr(cute.rank(mTargetLogit.shape) == 2):
-                # Use Int64 for indexing to deal with large tensors
-                mTargetLogit_off = utils.domain_offset_i64((row, 0), mTargetLogit)
-                target_logit = Float32(mTargetLogit_off[0, target])
+                target_logit = Float32(mTargetLogit[row, target])
             else:
                 assert cute.rank(mTargetLogit.shape) == 1
                 target_logit = Float32(mTargetLogit[row])
 
-        threads_per_row = tv_layout.shape[0][0]
         if const_expr(not self.online_softmax):
             max_x = row_reduce(
                 x,
@@ -237,21 +218,16 @@ class CrossEntropy(ReductionBase):
             # Compute probabilities: exp(x) / sum(exp(x))
             # If ignored, gradient should be zero
             denom_inv = (
-                1.0 / denom
+                # 1.0 / denom
+                cute.arch.rcp_approx(denom)
                 if not (denom == 0.0 or denom != denom or should_ignore)
                 else Float32.zero
             )
             probs = exp_x * denom_inv
-            mdX_off = utils.domain_offset_i64((bidx * tiler_mn[0], 0), mdX)
-            gdX = cute.local_tile(mdX_off, tiler_mn, (0, cluster_y))
-            # Setup copy atom for storing gradient
-            copy_atom_store = cute.make_copy_atom(
-                cute.nvgpu.CopyUniversalOp(), mdX.element_type, num_bits_per_copy=num_copy_bits_X
-            )
-            thr_copy_dX = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx)
-            tXgdX = thr_copy_dX.partition_D(gdX)
+            gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y))
+            tXgdX = thr_copy.partition_D(gdX)
             tXrdX = cute.make_fragment_like(tXgdX)
-            tXcFull = thr_copy_X.partition_S(cX)
+            tXcFull = thr_copy.partition_S(cX)
             # Compute gradient: probs for all classes, (probs - 1) for target class
             # If ignored, gradient is already zero
             tXrdX_f32 = cute.make_fragment_like(tXrX, Float32)
@@ -260,13 +236,8 @@ class CrossEntropy(ReductionBase):
                 for i in cutlass.range(cute.size(tXrX), unroll_full=True):
                     tXrdX_f32[i] = tXrdX_f32[i] if tXcFull[i][1] != target else tXrdX_f32[i] - 1.0
             tXrdX.store(tXrdX_f32.load().to(tXrdX.element_type))
-            tXpdX = (
-                utils.predicate_k(thr_copy_dX.partition_S(cX), limit=shape[1])
-                if not is_even_N
-                else None
-            )
             if row < shape[0]:
-                cute.copy(copy_atom_store, tXrdX, tXgdX, pred=tXpdX)
+                copy(tXrdX, tXgdX)
 
 
 @torch.library.custom_op("quack::cross_entropy_fwd_out", mutates_args={"loss", "lse", "dx"})
@@ -296,77 +267,61 @@ def cross_entropy_fwd_out(
     """
     assert x.dim() == 2, "Input must be 2D"
     assert target.dim() == 1, "Target must be 1D"
-    assert x.shape[0] == target.shape[0], "Batch dimensions must match"
     assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
     assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
     if target_logit is not None:
-        assert target_logit.shape[0] == x.shape[0]
         assert target_logit.is_cuda, "Target logits must be on CUDA device"
         assert target_logit.dtype in [torch.float16, torch.bfloat16, torch.float32]
     if dx is not None:
-        assert dx.shape == x.shape, "dx must have same shape as x"
         assert dx.is_cuda, "dx must be on CUDA device"
-        assert dx.dtype == x.dtype, "dx must have same dtype as x"
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
-    convert_from_dlpack = lambda tensor: (
-        from_dlpack(tensor.detach(), assumed_align=16).mark_compact_shape_dynamic(
-            mode=0, stride_order=(0, 1)
-        )
+    target_dtype = torch2cute_dtype_map[target.dtype]
+    target_logit_dtype = (
+        torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None
     )
-    x_tensor = convert_from_dlpack(x)
-    loss_tensor = from_dlpack(loss.detach(), assumed_align=4).mark_layout_dynamic()
-    lse_tensor = (
-        from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic()
-        if lse is not None
-        else None
-    )
-    target_tensor = from_dlpack(target.detach(), assumed_align=8).mark_layout_dynamic()
-    target_logit_tensor = (
-        from_dlpack(target_logit.detach(), assumed_align=4).mark_layout_dynamic(
-            leading_dim=target_logit.ndim - 1
-        )
-        if target_logit is not None
-        else None
-    )
-    dx_tensor = convert_from_dlpack(dx) if dx is not None else None
-    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-
     compile_key = (
         dtype,
+        target_dtype,
+        target_logit_dtype,
         N,
-        target_logit.dtype if target_logit is not None else None,
-        lse.dtype if lse is not None else None,
+        lse is not None,
         dx is not None,
-        loss.stride(),
-        lse.stride() if lse is not None else None,
-        target.stride(),
-        target_logit.stride(-1) if target_logit is not None else None,
     )
     if compile_key not in cross_entropy_fwd_out.compile_cache:
+        batch_sym = cute.sym_int()
+        div = math.gcd(128 // dtype.width, N)
+        x_cute = fake_tensor(dtype, (batch_sym, N), div)
+        dx_cute = fake_tensor(dtype, (batch_sym, N), div) if dx is not None else None
+        target_cute = fake_tensor(target_dtype, (batch_sym,))
+        if target_logit is not None:
+            if target_logit.ndim == 2:
+                target_logit_cute = fake_tensor(
+                    target_logit_dtype, (batch_sym, cute.sym_int()), div
+                )
+            else:
+                target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,))
+        else:
+            target_logit_cute = None
+        loss_cute = fake_tensor(Float32, (batch_sym,))
+        lse_cute = fake_tensor(Float32, (batch_sym,)) if lse is not None else None
         # If there's dx, it's faster to not use online softmax since we want the exp(x - max)
         cross_entropy_op = CrossEntropy(dtype, N, online_softmax=dx is None)
         cross_entropy_fwd_out.compile_cache[compile_key] = cute.compile(
             cross_entropy_op,
-            x_tensor,
-            target_tensor,
-            target_logit_tensor,
-            loss_tensor,
-            lse_tensor,
-            dx_tensor,
-            Int32(ignore_index),
-            stream,
+            x_cute,
+            target_cute,
+            target_logit_cute,
+            loss_cute,
+            lse_cute,
+            dx_cute,
+            Int32(0),  # ignore_index, just for compilation
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
         )
     cross_entropy_fwd_out.compile_cache[compile_key](
-        x_tensor,
-        target_tensor,
-        target_logit_tensor,
-        loss_tensor,
-        lse_tensor,
-        dx_tensor,
-        Int32(ignore_index),
-        stream,
+        x, target, target_logit, loss, lse, dx, Int32(ignore_index)
     )
 
 
@@ -404,35 +359,25 @@ class CrossEntropyBackward:
         self.N = N
         self.vecsize = 128 // dtype.width
 
-    def _calculate_threads_per_row(self):
+    def _threads_per_row(self):
         N = min(self.N, 16384)  # We split by blocks of 16k
-        return (
-            8
-            if N <= 64
-            else (
-                16
-                if N <= 128
-                else (32 if N <= 3072 else (64 if N <= 6144 else (128 if N <= 16384 else 256)))
-            )
-        )
+        for limit, threads in [(64, 8), (128, 16), (3072, 32), (6144, 64), (16384, 128)]:
+            if N <= limit:
+                return threads
+        return 256
 
-    def _get_tv_layout(self, num_copy_bits=128):
-        vecsize = num_copy_bits // self.dtype.width
+    def _get_tiled_copy(self, vecsize: int):
         assert self.N % vecsize == 0, f"Input N {self.N} is not divisible by vector size {vecsize}"
         N = min(self.N, 16384)
         num_threads = 128 if N <= 16384 else 256
-        threads_per_row = self._calculate_threads_per_row()
+        threads_per_row = self._threads_per_row()
         cols_per_block = num_threads // threads_per_row
         num_blocks_N = cute.ceil_div(N // vecsize, threads_per_row)
         tiler_mn = (cols_per_block, vecsize * num_blocks_N * threads_per_row)
-        tv_layout = cute.make_layout(
-            ((threads_per_row, cols_per_block), (vecsize, num_blocks_N)),
-            stride=(
-                (vecsize * cols_per_block, 1),
-                (cols_per_block, cols_per_block * vecsize * threads_per_row),
-            ),
+        tiled_copy = copy_utils.tiled_copy_2d(
+            self.dtype, threads_per_row, num_threads, num_copy_elems=vecsize
         )
-        return tiler_mn, tv_layout
+        return tiled_copy, tiler_mn, threads_per_row
 
     @cute.jit
     def __call__(
@@ -448,21 +393,24 @@ class CrossEntropyBackward:
         assert mX.element_type == self.dtype
         assert mdX.element_type == self.dtype
         # e.g. if self.N isn't divisible by 8 for bf16, we might use 64 bits (4 elements) copy
-        num_copy_bits = math.gcd(self.N, 128 // self.dtype.width) * self.dtype.width
-        tiler_mn, tv_layout = self._get_tv_layout(num_copy_bits=num_copy_bits)
-        num_threads = cute.size(tv_layout, mode=[0])
+        vecsize = math.gcd(self.N, 128 // self.dtype.width)
+        tiled_copy, tiler_mn, threads_per_row = self._get_tiled_copy(vecsize=vecsize)
+        num_threads = tiled_copy.size
         # (M,) -> (M, N) with stride 0 in the N dimension
         mDLoss, mTarget, mLSE = [
-            cute.make_tensor(
-                X.iterator, cute.append(X.layout, cute.make_layout((self.N,), stride=(0,)))
-            )
-            for X in (mDLoss, mTarget, mLSE)
+            layout_utils.expand(X, dim=1, size=self.N) for X in (mDLoss, mTarget, mLSE)
         ]
-        smem_size = cute.size_in_bytes(
-            mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0))
-        )
         self.kernel(
-            mX, mTarget, mDLoss, mdX, mLSE, ignore_index, mX.shape, tv_layout, tiler_mn
+            mX,
+            mTarget,
+            mDLoss,
+            mdX,
+            mLSE,
+            ignore_index,
+            mX.shape,
+            tiler_mn,
+            tiled_copy,
+            threads_per_row,
         ).launch(
             grid=[
                 cute.ceil_div(mX.shape[0], tiler_mn[0]),
@@ -470,7 +418,6 @@ class CrossEntropyBackward:
                 1,
             ],
             block=[num_threads, 1, 1],
-            smem=smem_size,
             stream=stream,
         )
 
@@ -484,52 +431,39 @@ class CrossEntropyBackward:
         mLSE: cute.Tensor,  # (M,)
         ignore_index: Int32,  # Index to ignore in gradient computation
         shape: cute.Shape,
-        tv_layout: cute.Layout,
         tiler_mn: cute.Shape,
+        tiled_copy: cute.TiledCopy,
+        threads_per_row: cutlass.Constexpr[int],
     ):
         tidx, _, _ = cute.arch.thread_idx()
         bidx, bidy, _ = cute.arch.block_idx()
 
         smem = cutlass.utils.SmemAllocator()
         sX = smem.allocate_tensor(
-            mX.element_type,
-            cute.make_ordered_layout(tiler_mn, order=(1, 0)),
-            byte_alignment=16,
+            mX.element_type, cute.make_ordered_layout(tiler_mn, order=(1, 0)), byte_alignment=16
         )
 
         idX = cute.make_identity_tensor(shape)
-        # We use domain_offset_i64 to deal with tensors larger than 2^31 elements
-        mX, mdX = [utils.domain_offset_i64((bidx * tiler_mn[0], 0), mT) for mT in (mX, mdX)]
-        gX, gdX = [cute.local_tile(mT, tiler_mn, (0, bidy)) for mT in (mX, mdX)]
-        cX = cute.local_tile(idX, tiler_mn, (bidx, bidy))
-
-        num_copy_elems_X = tv_layout.shape[1][0]
-        num_copy_bits_X = mX.element_type.width * num_copy_elems_X
-        copy_atom_load_X = cute.make_copy_atom(
-            cute.nvgpu.cpasync.CopyG2SOp(), gX.element_type, num_bits_per_copy=num_copy_bits_X
-        )
-        copy_atom_store_dX = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(), gdX.element_type, num_bits_per_copy=num_copy_bits_X
-        )
-        thr_copy_X = cute.make_tiled_copy(copy_atom_load_X, tv_layout, tiler_mn).get_slice(tidx)
-        thr_copy_dX = cute.make_tiled_copy(copy_atom_store_dX, tv_layout, tiler_mn).get_slice(tidx)
-
-        #### Partition to get thread view
-        tXgX = thr_copy_X.partition_S(gX)
-        tXsX = thr_copy_X.partition_S(sX)
-        tXcX = thr_copy_X.partition_S(cX)[(0, None), None, None]
-        tXcFull = thr_copy_X.partition_S(cX)
-        tXgdX = thr_copy_dX.partition_D(gdX)
-        # allocate fragments for gmem->rmem
+        gX, gdX, cX = [cute.local_tile(mT, tiler_mn, (bidx, bidy)) for mT in (mX, mdX, idX)]
+
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        tXgX = thr_copy.partition_S(gX)
+        tXsX = thr_copy.partition_D(sX)
+        tXcX = thr_copy.partition_S(cX)[(0, None), None, None]
+        tXcFull = thr_copy.partition_S(cX)
+        tXgdX = thr_copy.partition_D(gdX)
         tXrX, tXrdX = [cute.make_fragment_like(thr) for thr in (tXgX, tXgdX)]
 
         is_even_N = const_expr(shape[1] % tiler_mn[1] == 0)
-        row = tXcX[0][0]
         tXpX = (
-            utils.predicate_k(thr_copy_X.partition_S(cX), limit=shape[1]) if not is_even_N else None
+            None if is_even_N else copy_utils.predicate_k(thr_copy.partition_S(cX), limit=shape[1])
         )
+        copy = partial(copy_utils.copy, pred=tXpX)
+
+        row = tXcX[0][0]
         if row < shape[0]:
-            cute.copy(copy_atom_load_X, tXgX, tXsX, pred=tXpX)
+            copy(tXgX, tXsX, is_async=True)
         cute.arch.cp_async_commit_group()
         cute.arch.cp_async_wait_group(0)
         if const_expr(not is_even_N):
@@ -544,26 +478,22 @@ class CrossEntropyBackward:
             target = Int32(mTarget[row])
             should_ignore = Boolean(target == ignore_index)
             # Set dloss to 0 if this index should be ignored
-            dloss = Float32(mDLoss[row]) if not should_ignore else Float32.zero
+            if not should_ignore:
+                dloss = Float32(mDLoss[row])
             lse = Float32(mLSE[row])
 
         log2_e = math.log2(math.e)
         probs = cute.math.exp2(x * log2_e - (lse * log2_e), fastmath=True)
         prob_shifted = probs - 1.0
-        mask = cute.make_fragment_like(tXrX, cutlass.Boolean)
+        mask = cute.make_fragment_like(tXrX, Boolean)
         for i in cutlass.range(cute.size(tXcFull), unroll_full=True):
             mask[i] = tXcFull[i][1] == target
         grad = cute.where(mask.load(), prob_shifted, probs)
         grad = grad * dloss
 
         tXrdX.store(grad.to(tXrdX.element_type))
-        tXpdX = (
-            utils.predicate_k(thr_copy_dX.partition_S(cX), limit=shape[1])
-            if not is_even_N
-            else None
-        )
         if row < shape[0]:
-            cute.copy(copy_atom_store_dX, tXrdX, tXgdX, pred=tXpdX)
+            copy(tXrdX, tXgdX)
 
 
 def _cross_entropy_backward(
@@ -598,34 +528,28 @@ def _cross_entropy_backward(
 
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
-
-    convert_from_dlpack = lambda tensor: (
-        from_dlpack(tensor.detach(), assumed_align=16).mark_compact_shape_dynamic(
-            mode=0, stride_order=(0, 1)
-        )
-    )
-    x_tensor = convert_from_dlpack(x)
-    dx_tensor = convert_from_dlpack(dx)
-    dloss_tensor = from_dlpack(dloss.detach(), assumed_align=4).mark_layout_dynamic()
-    lse_tensor = from_dlpack(lse.detach(), assumed_align=4).mark_layout_dynamic()
-    target_tensor = from_dlpack(target.detach(), assumed_align=8).mark_layout_dynamic()
-    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
-
-    compile_key = (dtype, N, target.dtype, dloss.stride(), lse.stride(), target.stride())
+    target_dtype = torch2cute_dtype_map[target.dtype]
+    compile_key = (dtype, target_dtype, N)
     if compile_key not in _cross_entropy_backward.compile_cache:
+        batch_sym = cute.sym_int()
+        div = math.gcd(128 // dtype.width, N)
+        x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2
+        target_cute = fake_tensor(target_dtype, (batch_sym,))
+        dloss_cute, lse_cute = [fake_tensor(Float32, (batch_sym,))] * 2
         cross_entropy_backward_op = CrossEntropyBackward(dtype, N)
         _cross_entropy_backward.compile_cache[compile_key] = cute.compile(
             cross_entropy_backward_op,
-            x_tensor,
-            target_tensor,
-            dloss_tensor,
-            dx_tensor,
-            lse_tensor,
-            Int32(ignore_index),
-            stream,
+            x_cute,
+            target_cute,
+            dloss_cute,
+            dx_cute,
+            lse_cute,
+            Int32(0),  # ignore_index, just for compilation
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
         )
     _cross_entropy_backward.compile_cache[compile_key](
-        x_tensor, target_tensor, dloss_tensor, dx_tensor, lse_tensor, Int32(ignore_index), stream
+        x, target, dloss, dx, lse, Int32(ignore_index)
     )