# Copyright (c) 2025, Tri Dao. import math import operator from typing import Tuple from dataclasses import dataclass import cutlass import cutlass.cute as cute from cutlass import Float32 from .quack import layout_utils from . import utils from .quack.cute_dsl_utils import ParamsBase from .seqlen_info import SeqlenInfoQK @dataclass class Softmax(ParamsBase): scale_log2: Float32 num_rows: cutlass.Constexpr[int] row_max: cute.Tensor row_sum: cute.Tensor arch: cutlass.Constexpr[int] = 80 softmax_scale: Float32 | None = None @staticmethod def create( scale_log2: Float32, num_rows: cutlass.Constexpr[int], arch: cutlass.Constexpr[int] = 80, softmax_scale: Float32 | None = None, ): row_max = cute.make_rmem_tensor(num_rows, Float32) row_sum = cute.make_rmem_tensor(num_rows, Float32) return Softmax(scale_log2, num_rows, row_max, row_sum, arch, softmax_scale) def reset(self) -> None: self.row_max.fill(-Float32.inf) self.row_sum.fill(0.0) def _compute_row_max( self, acc_S_row: cute.TensorSSA, init_val: float | Float32 | None = None ) -> Float32: return utils.fmax_reduce(acc_S_row, init_val, arch=self.arch) def _compute_row_sum( self, acc_S_row_exp: cute.TensorSSA, init_val: float | Float32 | None = None ) -> Float32: return utils.fadd_reduce(acc_S_row_exp, init_val, arch=self.arch) @cute.jit def online_softmax( self, acc_S: cute.Tensor, is_first: cutlass.Constexpr[bool] = False, check_inf: cutlass.Constexpr[bool] = True, ) -> cute.Tensor: """Apply online softmax and return the row_scale to rescale O. :param acc_S: acc_S tensor :type acc_S: cute.Tensor :param is_first: is first n_block :type is_first: cutlass.Constexpr """ # Change acc_S to M,N layout view. acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) row_scale = cute.make_fragment_like(self.row_max, Float32) row_max = self.row_max row_sum = self.row_sum scale_log2 = self.scale_log2 arch = self.arch # Each iteration processes one row of acc_S for r in cutlass.range(cute.size(row_max), unroll_full=True): acc_S_row = acc_S_mn[r, None].load() # (n_block_size) row_max_cur = utils.fmax_reduce( acc_S_row, init_val=row_max[r] if cutlass.const_expr(not is_first) else None, arch=arch, ) row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4) # Update row_max before changing row_max_cur to safe value for -inf row_max_prev = row_max[r] row_max[r] = row_max_cur if cutlass.const_expr(check_inf): row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur if cutlass.const_expr(is_first): row_max_cur_scaled = row_max_cur * scale_log2 acc_S_row_exp = cute.math.exp2( acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True ) acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch) row_scale[r] = 1.0 else: row_max_cur_scaled = row_max_cur * scale_log2 acc_S_row_exp = cute.math.exp2( acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True ) # row_scale[r] = cute.math.exp2(row_max_prev * self.scale_log2 - row_max_cur_scaled) row_scale[r] = cute.math.exp2( (row_max_prev - row_max_cur) * scale_log2, fastmath=True ) acc_S_row_sum = utils.fadd_reduce( acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch ) row_sum[r] = acc_S_row_sum acc_S_mn[r, None].store(acc_S_row_exp) return row_scale @cute.jit def finalize( self, final_scale: Float32 = 1.0, sink_val: Float32 | cute.Tensor | None = None ) -> cute.Tensor: """Finalize the online softmax by computing the scale and logsumexp.""" if cutlass.const_expr(sink_val is not None and isinstance(sink_val, cute.Tensor)): assert cute.size(sink_val) == cute.size(self.row_sum) row_sum = self.row_sum row_max = self.row_max scale_log2 = self.scale_log2 # quad reduction for row_sum as we didn't do it during each iteration of online softmax row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4)) row_scale = cute.make_fragment_like(row_max, Float32) for r in cutlass.range(cute.size(row_sum), unroll_full=True): if cutlass.const_expr(sink_val is not None): sink_val_cur = sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r] LOG2_E = math.log2(math.e) row_sum[r] += cute.math.exp2( sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True ) # if row_sum is zero or nan, set acc_O_mn_row to 1.0 acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r] row_scale[r] = ( cute.arch.rcp_approx(row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0) ) * final_scale row_sum_cur = row_sum[r] LN2 = math.log(2.0) row_sum[r] = ( (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True)) * LN2 if not acc_O_mn_row_is_zero_or_nan else -Float32.inf ) return row_scale @cute.jit def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None: """Scale each row of acc_O by the given scale tensor. :param acc_O: input tensor :type acc_O: cute.Tensor :param row_scale: row_scale tensor :type row_scale: cute.Tensor """ acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O) assert cute.size(row_scale) == cute.size(acc_O_mn, mode=[0]) for r in cutlass.range(cute.size(row_scale), unroll_full=True): acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r]) @dataclass class SoftmaxSm100(Softmax): rescale_threshold: cutlass.Constexpr[float] = 0.0 @staticmethod def create( scale_log2: Float32, rescale_threshold: cutlass.Constexpr[float] = 0.0, softmax_scale: Float32 | None = None, ): num_rows = 1 arch = 100 row_max = cute.make_rmem_tensor(num_rows, Float32) row_sum = cute.make_rmem_tensor(num_rows, Float32) return SoftmaxSm100( scale_log2, num_rows, row_max, row_sum, arch, softmax_scale, rescale_threshold=rescale_threshold, ) @cute.jit def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]: if cutlass.const_expr(is_first): row_max_new = self._compute_row_max(acc_S_row) row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 acc_scale = 0.0 else: row_max_old = self.row_max[0] row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old) row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2 acc_scale = cute.math.exp2(acc_scale_, fastmath=True) if cutlass.const_expr(self.rescale_threshold > 0.0): if acc_scale_ >= -self.rescale_threshold: row_max_new = row_max_old row_max_safe = row_max_old acc_scale = 1.0 self.row_max[0] = row_max_new return row_max_safe, acc_scale def update_row_sum( self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False ) -> None: init_val = self.row_sum[0] * row_scale if cutlass.const_expr(not is_first) else None # self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=self.row_sum[0] * row_scale) self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val) # tmp = self._compute_row_sum(acc_S_row_exp) # self.row_sum[0] = self.row_sum[0] * row_scale + tmp @cute.jit def scale_subtract_rowmax( self, acc_S_row: cute.Tensor, row_max: Float32, ): assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" row_max_scaled = row_max * self.scale_log2 for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True): acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( (acc_S_row[i], acc_S_row[i + 1]), (self.scale_log2, self.scale_log2), (-row_max_scaled, -row_max_scaled), ) @cute.jit def apply_exp2_convert( self, acc_S_row: cute.Tensor, acc_S_row_converted: cute.Tensor, ex2_emu_freq: cutlass.Constexpr[int] = 0, ex2_emu_res: cutlass.Constexpr[int] = 4, ex2_emu_start_frg: cutlass.Constexpr[int] = 0, ): assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" frg_tile = 32 assert frg_tile % 2 == 0 frg_cnt = cute.size(acc_S_row) // frg_tile assert cute.size(acc_S_row) % frg_tile == 0 acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) acc_S_row_converted_frg = cute.logical_divide( acc_S_row_converted, cute.make_layout(frg_tile) ) for j in cutlass.range_constexpr(frg_cnt): for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) if cutlass.const_expr(ex2_emu_freq == 0): acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) else: if cutlass.const_expr( k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res or j >= frg_cnt - 1 or j < ex2_emu_start_frg ): acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) acc_S_row_frg[k + 1, j] = cute.math.exp2( acc_S_row_frg[k + 1, j], fastmath=True ) else: # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.e2e_asm2(acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]) acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.ex2_emulation_2( acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] ) acc_S_row_converted_frg[None, j].store( acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) ) @cute.jit def scale_apply_exp2_convert( self, acc_S_row: cute.Tensor, row_max: Float32, acc_S_row_converted: cute.Tensor, ): assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" minus_row_max_scaled = -row_max * self.scale_log2 for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( (acc_S_row[i], acc_S_row[i + 1]), (self.scale_log2, self.scale_log2), (minus_row_max_scaled, minus_row_max_scaled), ) # for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): # acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( # (acc_S_row[i], acc_S_row[i + 1]), # (self.scale_log2, self.scale_log2), # (minus_row_max_scaled, minus_row_max_scaled), # ) # acc_S_row[i] = cute.math.exp2(acc_S_row[i], fastmath=True) # acc_S_row[i + 1] = cute.math.exp2(acc_S_row[i + 1], fastmath=True) frg_tile = 32 assert frg_tile % 2 == 0 frg_cnt = cute.size(acc_S_row) // frg_tile assert cute.size(acc_S_row) % frg_tile == 0 acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) acc_S_row_converted_frg = cute.logical_divide( acc_S_row_converted, cute.make_layout(frg_tile) ) for j in cutlass.range_constexpr(frg_cnt): for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = ( # cute.arch.fma_packed_f32x2( # (acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]), # (self.scale_log2, self.scale_log2), # (minus_row_max_scaled, minus_row_max_scaled), # ) # ) # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) acc_S_row_converted_frg[None, j].store( acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) ) @cute.jit def floor_if_packed( q_idx, qhead_per_kvhead: cutlass.Constexpr[int], ) -> cute.Tensor: """Convert q_idx to packed format for Pack-GQA.""" if cutlass.const_expr(qhead_per_kvhead == 1): return q_idx return q_idx // qhead_per_kvhead @cute.jit def apply_score_mod_inner( score_tensor, index_tensor, score_mod: cutlass.Constexpr, batch_idx, head_idx, softmax_scale, vec_size: cutlass.Constexpr, qk_acc_dtype: cutlass.Constexpr, aux_tensors, fastdiv_mods, seqlen_info: SeqlenInfoQK, constant_q_idx: cutlass.Constexpr, qhead_per_kvhead: cutlass.Constexpr[int] = 1, transpose_indices: cutlass.Constexpr[bool] = False, ): """Shared implementation for applying score modification. Args: score_tensor: The scores to modify (acc_S for flash_fwd, tSrS_t2r for sm100) index_tensor: Index positions (tScS for flash_fwd, tScS_t2r for sm100) score_mod: The score modification function to apply batch_idx: Batch index head_idx: Head index softmax_scale: Scale to apply vec_size: Vector size for processing elements qk_acc_dtype: Data type for accumulator aux_tensors: Optional aux_tensors for FlexAttention fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping seqlen_info: Sequence length info constant_q_idx: If provided, use this constant for all q_idx values If None, compute q_idx per-element qhead_per_kvhead_packgqa: Pack-GQA replication factor. Divide q_idx by this when greater than 1 so score mods see logical heads. transpose_indices: If True, swap q_idx/kv_idx in index_tensor (for bwd kernel where S is transposed) """ # Index positions in the index_tensor tuple # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx if cutlass.const_expr(transpose_indices): q_idx_pos = cutlass.const_expr(1) kv_idx_pos = cutlass.const_expr(0) else: q_idx_pos = cutlass.const_expr(0) kv_idx_pos = cutlass.const_expr(1) n_vals = cutlass.const_expr(cute.size(score_tensor.shape)) score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) # SSA values for batch (constant across all elements) batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) # Handle q_idx based on whether it's constant q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) # For Pack-GQA with non-constant q_idx, we need per-element head indices # since a thread my process multiple query head indices if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): for j in cutlass.range(vec_size, unroll_full=True): score_vec[j] = score_tensor[i + j] * softmax_scale # Extract head offset from packed q_idx for Pack-GQA if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): q_idx_packed = index_tensor[i + j][q_idx_pos] # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead) q_idx_logical = q_idx_packed // qhead_per_kvhead head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset # If we will do loads we mod, in order to not read OOB if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None): if cutlass.const_expr(constant_q_idx is None): seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods q_idx_floored = floor_if_packed( index_tensor[i + j][q_idx_pos], qhead_per_kvhead ) _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) q_idx_vec[j] = q_idx_wrapped else: _, seqlen_k_divmod = fastdiv_mods _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) kv_idx_vec[j] = kv_idx_wrapped else: # No bounds checking - direct indexing if constant_q_idx is None: q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] # Convert to SSA for score_mod call score_ssa = score_vec.load() kv_idx_ssa = kv_idx_vec.load() if cutlass.const_expr(constant_q_idx is None): q_idx_ssa = q_idx_vec.load() else: # NB we do not apply Pack-GQA division here, as constant_q_idx is assumed to already be logical q_idx_const = constant_q_idx q_idx_ssa = utils.scalar_to_ssa(q_idx_const, cutlass.Int32).broadcast_to((vec_size,)) # Compute head_idx_ssa: per-element for Pack-GQA with non-constant q_idx, constant otherwise if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): head_idx_ssa = head_idx_vec.load() else: head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) aux_args = [] if cutlass.const_expr(aux_tensors is not None): aux_args = aux_tensors post_mod_scores = score_mod( score_ssa, batch_idx_ssa, head_idx_ssa, q_idx=q_idx_ssa, kv_idx=kv_idx_ssa, seqlen_info=seqlen_info, aux_tensors=aux_args, ) # Write back modified scores score_vec.store(post_mod_scores) for j in cutlass.range(vec_size, unroll_full=True): score_tensor[i + j] = score_vec[j] @cute.jit def apply_score_mod_bwd_inner( grad_tensor, score_tensor, index_tensor, score_mod_bwd: cutlass.Constexpr, batch_idx, head_idx, softmax_scale, vec_size: cutlass.Constexpr, qk_acc_dtype: cutlass.Constexpr, aux_tensors, fastdiv_mods, seqlen_info, constant_q_idx: cutlass.Constexpr, qhead_per_kvhead: cutlass.Constexpr[int] = 1, transpose_indices: cutlass.Constexpr[bool] = False, ): """Apply backward score modification (joint graph). Args: grad_tensor: in/out: dlogits rewritten in-place with d(scaled_scores) score_tensor: pre-mod scores (unscaled QK tile), scaled by softmax_scale internally index_tensor: Index positions (same as forward) score_mod_bwd: The backward score modification function (joint graph) batch_idx: Batch index head_idx: Head index softmax_scale: Scale to apply to score_tensor vec_size: Vector size for processing elements qk_acc_dtype: Data type for accumulator aux_tensors: Optional aux_tensors for FlexAttention fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping seqlen_info: Sequence length info constant_q_idx: If provided, use this constant for all q_idx values qhead_per_kvhead: Pack-GQA replication factor transpose_indices: If True, swap q_idx/kv_idx in index_tensor """ # Index positions in the index_tensor tuple # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx if cutlass.const_expr(transpose_indices): q_idx_pos = cutlass.const_expr(1) kv_idx_pos = cutlass.const_expr(0) else: q_idx_pos = cutlass.const_expr(0) kv_idx_pos = cutlass.const_expr(1) n_vals = cutlass.const_expr(cute.size(grad_tensor.shape)) grad_vec = cute.make_fragment(vec_size, qk_acc_dtype) score_vec = cute.make_fragment(vec_size, qk_acc_dtype) kv_idx_vec = cute.make_fragment(vec_size, cutlass.Int32) batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) q_idx_vec = cute.make_fragment(vec_size, cutlass.Int32) # For Pack-GQA with non-constant q_idx, we need per-element head indices if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): head_idx_vec = cute.make_fragment(vec_size, cutlass.Int32) for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): for j in cutlass.range(vec_size, unroll_full=True): grad_vec[j] = grad_tensor[i + j] # Scale score so joint graph sees same value as forward score_mod score_vec[j] = score_tensor[i + j] * softmax_scale if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): q_idx_packed = index_tensor[i + j][q_idx_pos] q_idx_logical = q_idx_packed // qhead_per_kvhead head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset if cutlass.const_expr(aux_tensors is not None and fastdiv_mods is not None): if cutlass.const_expr(constant_q_idx is None): seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods q_idx_floored = floor_if_packed( index_tensor[i + j][q_idx_pos], qhead_per_kvhead ) _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) q_idx_vec[j] = q_idx_wrapped else: _, seqlen_k_divmod = fastdiv_mods _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) kv_idx_vec[j] = kv_idx_wrapped else: # No bounds checking - direct indexing if constant_q_idx is None: q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] grad_ssa = grad_vec.load() score_ssa = score_vec.load() kv_idx_ssa = kv_idx_vec.load() if cutlass.const_expr(constant_q_idx is None): q_idx_ssa = q_idx_vec.load() else: q_idx_ssa = utils.scalar_to_ssa(constant_q_idx, cutlass.Int32).broadcast_to((vec_size,)) if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): head_idx_ssa = head_idx_vec.load() else: head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) aux_args = [] if cutlass.const_expr(aux_tensors is not None): aux_args = aux_tensors grad_out_ssa = score_mod_bwd( grad_ssa, score_ssa, batch_idx_ssa, head_idx_ssa, q_idx=q_idx_ssa, kv_idx=kv_idx_ssa, seqlen_info=seqlen_info, aux_tensors=aux_args, ) grad_vec.store(grad_out_ssa) for j in cutlass.range(vec_size, unroll_full=True): grad_tensor[i + j] = grad_vec[j]