File size: 11,075 Bytes
8c8128e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 | from typing import Optional
from dataclasses import dataclass
import cutlass
import cutlass.cute as cute
from cutlass import Int32, const_expr
from fa4_cute_runtime.quack import copy_utils
"""
This consolidates all the info related to sequence length. This is so that we can do all
the gmem reads once at the beginning of each tile, rather than having to repeat these reads
to compute various things like n_block_min, n_block_max, etc.
"""
@dataclass(frozen=True)
class SeqlenInfo:
offset: Int32
offset_padded: Int32
seqlen: Int32
has_cu_seqlens: cutlass.Constexpr[bool] = False
@staticmethod
def create(
batch_idx: Int32,
seqlen_static: Int32,
cu_seqlens: Optional[cute.Tensor] = None,
seqused: Optional[cute.Tensor] = None,
tile: cutlass.Constexpr[int] = 128,
):
offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx]
offset_padded = (
0
if const_expr(cu_seqlens is None)
# Add divby so that the compiler knows the alignment when moving by offset_padded
else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile)
)
if const_expr(seqused is not None):
seqlen = seqused[batch_idx]
elif const_expr(cu_seqlens is not None):
seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
else:
seqlen = seqlen_static
return SeqlenInfo(offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None)
def offset_batch(
self,
mT: cute.Tensor,
batch_idx: Int32,
dim: int,
padded: cutlass.Constexpr[bool] = False,
multiple: int = 1,
) -> cute.Tensor:
"""Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0."""
if const_expr(not self.has_cu_seqlens):
idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim)
return mT[idx]
else:
off = multiple * (self.offset if const_expr(not padded) else self.offset_padded)
offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off)
idx = (offset,) + (None,) * (cute.rank(mT) - 1)
return cute.domain_offset(idx, mT)
@dataclass(frozen=True)
class SeqlenInfoQK:
offset_q: Int32
offset_k: Int32
padded_offset_q: Int32
padded_offset_k: Int32
seqlen_q: Int32
seqlen_k: Int32
has_cu_seqlens_q: cutlass.Constexpr[bool]
has_cu_seqlens_k: cutlass.Constexpr[bool]
has_seqused_q: cutlass.Constexpr[bool]
has_seqused_k: cutlass.Constexpr[bool]
@staticmethod
def create(
batch_idx: Int32,
seqlen_q_static: Int32,
seqlen_k_static: Int32,
mCuSeqlensQ: Optional[cute.Tensor] = None,
mCuSeqlensK: Optional[cute.Tensor] = None,
mSeqUsedQ: Optional[cute.Tensor] = None,
mSeqUsedK: Optional[cute.Tensor] = None,
tile_m: cutlass.Constexpr[Int32] = 128,
tile_n: cutlass.Constexpr[Int32] = 128,
):
offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx]
padded_offset_q = (
0
if const_expr(mCuSeqlensQ is None)
else cute.assume((offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m)
)
padded_offset_k = (
0
if const_expr(mCuSeqlensK is None)
else cute.assume((offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n)
)
if const_expr(mSeqUsedQ is not None):
seqlen_q = mSeqUsedQ[batch_idx]
else:
seqlen_q = (
seqlen_q_static
if const_expr(mCuSeqlensQ is None)
else mCuSeqlensQ[batch_idx + 1] - offset_q
)
if const_expr(mSeqUsedK is not None):
seqlen_k = mSeqUsedK[batch_idx]
else:
seqlen_k = (
seqlen_k_static
if const_expr(mCuSeqlensK is None)
else mCuSeqlensK[batch_idx + 1] - offset_k
)
return SeqlenInfoQK(
offset_q,
offset_k,
padded_offset_q,
padded_offset_k,
seqlen_q,
seqlen_k,
has_cu_seqlens_q=mCuSeqlensQ is not None,
has_cu_seqlens_k=mCuSeqlensK is not None,
has_seqused_q=mSeqUsedQ is not None,
has_seqused_k=mSeqUsedK is not None,
)
def offset_batch_Q(
self,
mQ: cute.Tensor,
batch_idx: Int32,
dim: int,
padded: cutlass.Constexpr[bool] = False,
ragged: cutlass.Constexpr[bool] = False,
) -> cute.Tensor:
"""Seqlen must be the first dimension of mQ"""
if const_expr(not ragged):
if const_expr(not self.has_cu_seqlens_q):
idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
return mQ[idx]
else:
offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
offset_q = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (None, offset_q)
idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1)
return cute.domain_offset(idx, mQ)
else:
if const_expr(not self.has_cu_seqlens_q):
offset_q = 0
idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim)
mQ = mQ[idx]
else:
offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q
if const_expr(cute.rank(mQ.shape[0]) == 1):
return copy_utils.offset_ragged_tensor(
mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True
)
else: # PackGQA
assert cute.rank(mQ.shape[0]) == 2
# Unpack before calling offset_ragged_tensor, then pack
idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1)
mQ = mQ[idx]
mQ = copy_utils.offset_ragged_tensor(
mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True
)
return cute.group_modes(mQ, 0, 2)
def offset_batch_K(
self,
mK: cute.Tensor,
batch_idx: Int32,
dim: int,
padded: cutlass.Constexpr[bool] = False,
ragged: cutlass.Constexpr[bool] = False,
multiple: int = 1,
) -> cute.Tensor:
"""Seqlen must be the first dimension of mK"""
if const_expr(not ragged):
if const_expr(not self.has_cu_seqlens_k):
idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
return mK[idx]
else:
offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
offset_k *= multiple
idx = (offset_k,) + (None,) * (cute.rank(mK) - 1)
return cute.domain_offset(idx, mK)
else:
if const_expr(not self.has_cu_seqlens_k):
offset_k = 0
idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim)
mK = mK[idx]
else:
offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k
offset_k *= multiple
return copy_utils.offset_ragged_tensor(
mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True
)
@dataclass(frozen=True)
class SeqlenInfoQKNewK:
"""Sequence length info for append-KV with left-padding and new K support.
Extends SeqlenInfoQK with:
- leftpad_k: left padding for K (tokens to skip at the start of the KV cache)
- offset_k_new: offset into the new K tensor
- seqlen_k_og: original K length (before appending new K), excluding leftpad
- seqlen_k_new: length of new K to append
- seqlen_k: total K length (seqlen_k_og + seqlen_k_new)
- seqlen_rotary: position for rotary embedding computation
"""
leftpad_k: Int32
offset_q: Int32
offset_k: Int32
offset_k_new: Int32
seqlen_q: Int32
seqlen_k_og: Int32
seqlen_k_new: Int32
seqlen_k: Int32
seqlen_rotary: Int32
@staticmethod
def create(
batch_idx: Int32,
seqlen_q_static: Int32,
seqlen_k_static: Int32,
shape_K_new_0: Int32,
mCuSeqlensQ: Optional[cute.Tensor] = None,
mCuSeqlensK: Optional[cute.Tensor] = None,
mCuSeqlensKNew: Optional[cute.Tensor] = None,
mSeqUsedQ: Optional[cute.Tensor] = None,
mSeqUsedK: Optional[cute.Tensor] = None,
mLeftpadK: Optional[cute.Tensor] = None,
mSeqlensRotary: Optional[cute.Tensor] = None,
):
leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx]
offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx]
if const_expr(mCuSeqlensK is not None):
offset_k = mCuSeqlensK[batch_idx] + leftpad_k
else:
offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0
offset_k_new = 0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx]
# seqlen_q
if const_expr(mSeqUsedQ is not None):
seqlen_q = mSeqUsedQ[batch_idx]
elif const_expr(mCuSeqlensQ is not None):
seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx]
else:
seqlen_q = seqlen_q_static
# seqlen_k_og: original K length (excluding leftpad)
if const_expr(mSeqUsedK is not None):
seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k
elif const_expr(mCuSeqlensK is not None):
seqlen_k_og = mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k
else:
seqlen_k_og = (
seqlen_k_static - leftpad_k
if const_expr(mCuSeqlensQ is not None)
else seqlen_k_static
)
# seqlen_k_new
if const_expr(mCuSeqlensKNew is None):
seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0
else:
seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx]
seqlen_k = seqlen_k_og if const_expr(mCuSeqlensQ is None) else seqlen_k_og + seqlen_k_new
# seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided
if const_expr(mSeqlensRotary is not None):
seqlen_rotary = mSeqlensRotary[batch_idx]
else:
seqlen_rotary = seqlen_k_og + leftpad_k
return SeqlenInfoQKNewK(
leftpad_k,
offset_q,
offset_k,
offset_k_new,
seqlen_q,
seqlen_k_og,
seqlen_k_new,
seqlen_k,
seqlen_rotary,
)
|