iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[GDN] Improve kkt kernel of CuteDSL prefill backend (#46346)

Browse Source 
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
This commit is contained in:
Thien Tran
2026-06-29 18:34:18 -07:00
committed by GitHub
parent cda05ee8c4
commit 43916891b2
6 changed files with 559 additions and 479 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/cute_utils/__init__.py
+35 -21
View File
@@ -95,40 +95,54 @@ def mma_bf16(
return cute.TensorSSA(vec, 4, Float32)
@dsl_user_op
def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32:
out = llvm.inline_asm(
T.i32(),
[a.ir_value(loc=loc, ip=ip)],
"abs.bf16x2 $0, $1;",
f"{asm}.bf16x2 $0, $1;",
"=r,r",
has_side_effects=False,
is_align_stack=False,
loc=loc,
ip=ip,
)
return Uint32(out)
def _bf16x2_binary(asm: str, a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
out = llvm.inline_asm(
T.i32(),
[a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
f"{asm}.bf16x2 $0, $1, $2;",
"=r,r,r",
has_side_effects=False,
is_align_stack=False,
loc=loc,
ip=ip,
)
return Uint32(out)
@dsl_user_op
def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
return _bf16x2_unary("abs", a, loc=loc, ip=ip)
@dsl_user_op
def _bf16x2_neg(a: Uint32, *, loc=None, ip=None) -> Uint32:
return _bf16x2_unary("neg", a, loc=loc, ip=ip)
@dsl_user_op
def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
out = llvm.inline_asm(
T.i32(),
[a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
"max.bf16x2 $0, $1, $2;",
"=r,r,r",
has_side_effects=False,
is_align_stack=False,
)
return Uint32(out)
return _bf16x2_binary("max", a, b, loc=loc, ip=ip)
@dsl_user_op
def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
out = llvm.inline_asm(
T.i32(),
[a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
"mul.rn.bf16x2 $0, $1, $2;",
"=r,r,r",
has_side_effects=False,
is_align_stack=False,
)
return Uint32(out)
return _bf16x2_binary("mul.rn", a, b, loc=loc, ip=ip)
@dsl_user_op
def _bf16x2_sub(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
return _bf16x2_binary("sub.rn", a, b, loc=loc, ip=ip)
vllm/cute_utils/_tcgen05.py
+34 -30
View File
@@ -90,17 +90,18 @@ def mma_f16(
loc=None,
ip=None,
) -> None:
nvvm.tcgen05_mma(
nvvm.Tcgen05MMAKind.F16,
NVVM_CTA_GROUP_MAP[cta_group],
_make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
Uint64(a_desc).ir_value(loc=loc, ip=ip),
Uint64(b_desc).ir_value(loc=loc, ip=ip),
Int32(idesc).ir_value(loc=loc, ip=ip),
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
with cute.arch.elect_one():
nvvm.tcgen05_mma(
nvvm.Tcgen05MMAKind.F16,
NVVM_CTA_GROUP_MAP[cta_group],
_make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
Uint64(a_desc).ir_value(loc=loc, ip=ip),
Uint64(b_desc).ir_value(loc=loc, ip=ip),
Int32(idesc).ir_value(loc=loc, ip=ip),
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
@dsl_user_op
@@ -115,17 +116,18 @@ def mma_ts_f16(
loc=None,
ip=None,
) -> None:
nvvm.tcgen05_mma(
nvvm.Tcgen05MMAKind.F16,
NVVM_CTA_GROUP_MAP[cta_group],
_make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
_make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
Uint64(b_desc).ir_value(loc=loc, ip=ip),
Int32(idesc).ir_value(loc=loc, ip=ip),
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
with cute.arch.elect_one():
nvvm.tcgen05_mma(
nvvm.Tcgen05MMAKind.F16,
NVVM_CTA_GROUP_MAP[cta_group],
_make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
_make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
Uint64(b_desc).ir_value(loc=loc, ip=ip),
Int32(idesc).ir_value(loc=loc, ip=ip),
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
@dsl_user_op
@@ -133,15 +135,17 @@ def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
group = NVVM_CTA_GROUP_MAP[cta_group]
if cutlass.const_expr(cta_mask is not None):
nvvm.tcgen05_commit_arrive(
mbar_llvm,
multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
group=group,
loc=loc,
ip=ip,
)
with cute.arch.elect_one():
nvvm.tcgen05_commit_arrive(
mbar_llvm,
multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
group=group,
loc=loc,
ip=ip,
)
else:
nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
with cute.arch.elect_one():
nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
@dsl_user_op
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
+22 -27
View File
@@ -174,31 +174,31 @@ def chunk_gated_delta_rule_cutedsl(
When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of
that buffer.
"""
q_3d = q.squeeze(0)
k_3d = k.squeeze(0)
v_3d = v.squeeze(0)
g_2d = g.squeeze(0)
beta_2d = beta.squeeze(0)
q = q.squeeze(0)
k = k.squeeze(0)
v = v.squeeze(0)
g = g.squeeze(0)
beta = beta.squeeze(0)
_, _, head_k_dim = k_3d.shape
_, num_v_heads, head_v_dim = v_3d.shape
_, _, K_dim = k.shape
_, num_v_heads, V_dim = v.shape
chunk_size = 64
upper_bound_chunks = chunk_indices.shape[0]
pad_t = upper_bound_chunks * chunk_size
total_chunks_ptr = chunk_offsets[-1:]
g_cu = torch.empty_like(g_2d, dtype=torch.float32)
u = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
w = q_3d.new_empty(pad_t, num_v_heads, head_k_dim)
g_cu = torch.empty_like(g, dtype=torch.float32)
u = q.new_empty(pad_t, num_v_heads, V_dim)
w = q.new_empty(pad_t, num_v_heads, K_dim)
num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
kkt_inv_uw_cutedsl(
k_3d,
v_3d,
k,
v,
u,
w,
g_2d,
beta_2d,
g,
beta,
g_cu,
cu_seqlens,
chunk_indices,
@@ -206,16 +206,11 @@ def chunk_gated_delta_rule_cutedsl(
num_sms=num_sms,
)
h = k_3d.new_empty(
upper_bound_chunks,
num_v_heads,
head_v_dim,
head_k_dim,
)
v_new = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
h = k.new_empty(upper_bound_chunks, num_v_heads, V_dim, K_dim)
v_new = q.new_empty(pad_t, num_v_heads, V_dim)
final_state = torch.empty_like(initial_state)
h_cutedsl(
k_3d,
k,
u,
w,
v_new,
@@ -227,12 +222,12 @@ def chunk_gated_delta_rule_cutedsl(
chunk_offsets,
)
output = core_attn_out if core_attn_out is not None else torch.empty_like(v_3d)
scale = head_k_dim**-0.5
output = core_attn_out if core_attn_out is not None else torch.empty_like(v)
scale = K_dim**-0.5
o_cutedsl(
q_3d,
k_3d,
v_new.view(upper_bound_chunks, chunk_size, num_v_heads, head_v_dim),
q,
k,
v_new,
h,
g_cu,
output,
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
+89 -90
View File
@@ -66,13 +66,12 @@ class Sm100ChunkHKernel:
stride=(64, 0, (1, self.BT * 64), self.BT * dim),
)
slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
atom, tma_tensor = cpasync.make_tiled_tma_atom(
return cpasync.make_tiled_tma_atom(
op,
cute.logical_divide(tensor, (None, None, 64)),
slayout,
cta_tiler=(self.BT, 1, dim),
)
return atom, tma_tensor, slayout
@cute.jit
def _make_h_tma_args(self, tensor: cute.Tensor, op: cpasync.TmaCopyOp):
@@ -84,13 +83,12 @@ class Sm100ChunkHKernel:
stride=(0, 0, num_elems, (1, self.V_dim * num_elems)),
)
slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
atom, tma_tensor = cpasync.make_tiled_tma_atom(
return cpasync.make_tiled_tma_atom(
op,
cute.logical_divide(tensor, (None, None, None, num_elems)),
slayout,
cta_tiler=(1, 1, self.V_dim, self.K_dim),
)
return atom, tma_tensor, slayout
@cute.jit
def __call__(
@@ -110,39 +108,39 @@ class Sm100ChunkHKernel:
tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
K_args = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages)
V_args = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages)
W_args = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages)
V_new_args = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1)
H0_args = self._make_h_tma_args(h0, tma_g2s)
HT_args = self._make_h_tma_args(ht, tma_s2g)
H_args = self._make_h_tma_args(h, tma_s2g)
K_tma = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages)
V_tma = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages)
W_tma = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages)
V_new_tma = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1)
H0_tma = self._make_h_tma_args(h0, tma_g2s)
HT_tma = self._make_h_tma_args(ht, tma_s2g)
H_tma = self._make_h_tma_args(h, tma_s2g)
grid = (self.Hv, h0.shape[0], 1)
block = (self.num_warps * 32, 1, 1)
self.kernel(
K_args,
V_args,
W_args,
V_new_args,
H0_args,
HT_args,
H_args,
K_tma,
V_tma,
W_tma,
V_new_tma,
H0_tma,
HT_tma,
H_tma,
g_cu,
cu_seqlens,
chunk_offsets,
).launch(grid=grid, block=block, stream=stream)
).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
@cute.kernel
def kernel(
self,
K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
V_new_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
H0_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
HT_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
K_tma: cpasync.TmaInfo,
V_tma: cpasync.TmaInfo,
W_tma: cpasync.TmaInfo,
V_new_tma: cpasync.TmaInfo,
H0_tma: cpasync.TmaInfo,
HT_tma: cpasync.TmaInfo,
H_tma: cpasync.TmaInfo,
g_cu: cute.Tensor,
cu_seqlens: cute.Tensor,
chunk_offsets: cute.Tensor,
@@ -158,14 +156,6 @@ class Sm100ChunkHKernel:
num_stages = self.num_stages
is_f32 = self.h_dtype == Float32
K_tma_atom, tmaK, sK_layout = K_args
V_tma_atom, tmaV, sV_layout = V_args
W_tma_atom, tmaW, sW_layout = W_args
V_new_tma_atom, tmaV_new, sV_new_layout = V_new_args
H0_tma_atom, tmaH0, sH0_layout = H0_args
HT_tma_atom, tmaHT, _ = HT_args
H_tma_atom, tmaH, sH_layout = H_args
def allocate_tensor(smem, dtype, layout):
return smem.allocate_tensor(
dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
@@ -174,12 +164,14 @@ class Sm100ChunkHKernel:
smem = cutlass.utils.SmemAllocator()
# remove size=1 modes
sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
sH0 = allocate_tensor(smem, self.h_dtype, sH0_layout)[0, 0, None, None]
sH = allocate_tensor(smem, BFloat16, sH_layout)[0, 0, None, None]
sV_new = allocate_tensor(smem, BFloat16, sV_new_layout)[None, 0, None, 0]
sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
sH0 = allocate_tensor(smem, self.h_dtype, H0_tma.smem_layout)[0, 0, None, None]
sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None]
sV_new = allocate_tensor(smem, BFloat16, V_new_tma.smem_layout)[
None, 0, None, 0
]
s_v_scale = smem.allocate_array(Float32, BT)
tma_mbar = smem.allocate_array(Int64, num_stages)
@@ -206,13 +198,13 @@ class Sm100ChunkHKernel:
cute.arch.mbarrier_init(h0_mbar, 1)
cute.arch.mbarrier_init_fence()
elif warp_id == 1:
cpasync.prefetch_descriptor(H0_tma_atom)
cpasync.prefetch_descriptor(W_tma_atom)
cpasync.prefetch_descriptor(V_tma_atom)
cpasync.prefetch_descriptor(K_tma_atom)
cpasync.prefetch_descriptor(HT_tma_atom)
cpasync.prefetch_descriptor(H_tma_atom)
cpasync.prefetch_descriptor(V_new_tma_atom)
cpasync.prefetch_descriptor(H0_tma.atom)
cpasync.prefetch_descriptor(W_tma.atom)
cpasync.prefetch_descriptor(V_tma.atom)
cpasync.prefetch_descriptor(K_tma.atom)
cpasync.prefetch_descriptor(HT_tma.atom)
cpasync.prefetch_descriptor(H_tma.atom)
cpasync.prefetch_descriptor(V_new_tma.atom)
cute.arch.sync_threads()
bos = cu_seqlens[seq_id]
@@ -233,14 +225,21 @@ class Sm100ChunkHKernel:
H0_size = V_dim * K_dim * self.h_dtype.width // 8
cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
simple_tma_copy(
H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
H0_tma.atom,
H0_tma.tma_tensor[seq_id, head_id, None, None],
sH0,
h0_mbar,
)
# shape: ((BT, num_BT_tiles), (64, 2))
gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
gV_tiles = cute.logical_divide(tmaV[None, head_id, None], (BT, None))
gW_tiles = cute.logical_divide(
W_tma.tma_tensor[None, head_id, None], (BT, None)
)
gV_tiles = cute.logical_divide(
V_tma.tma_tensor[None, head_id, None], (BT, None)
)
gK_tiles = cute.logical_divide(
cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
cute.domain_offset((bos, 0), K_tma.tma_tensor[None, k_head_id, None]),
(BT, None),
)
@@ -258,12 +257,12 @@ class Sm100ChunkHKernel:
STAGE_SIZE = BT * (K_dim + V_dim + K_dim) * 2
cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
simple_tma_copy(
W_tma_atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
W_tma.atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
)
simple_tma_copy(
V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
)
simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
@@ -295,13 +294,12 @@ class Sm100ChunkHKernel:
cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4)
wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True)
_tcgen05.commit(wh_done_mbar + stage_id)
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4)
wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True)
_tcgen05.commit(wh_done_mbar + stage_id)
##### 2nd MMA: H_new = H + V_new.T @ K #####
Kaddr0 = sK[None, None, stage_id].iterator.toint()
@@ -310,12 +308,11 @@ class Sm100ChunkHKernel:
cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for k in cutlass.range_constexpr(BT // 16):
vtmem0 = v_tmem_base + k * 8
kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True)
_tcgen05.commit(vk_done_mbar + stage_id)
for k in cutlass.range_constexpr(BT // 16):
vtmem0 = v_tmem_base + k * 8
kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True)
_tcgen05.commit(vk_done_mbar + stage_id)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
@@ -331,13 +328,12 @@ class Sm100ChunkHKernel:
cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
htmem = h_tmem_base + i * 32 + j * 8
wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True)
_tcgen05.commit(wh_done_mbar + stage_id)
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
htmem = h_tmem_base + i * 32 + j * 8
wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True)
_tcgen05.commit(wh_done_mbar + stage_id)
##### 2nd MMA: H_new = H + V_new.T @ K #####
Kaddr = sK[None, None, stage_id].iterator.toint()
@@ -346,12 +342,11 @@ class Sm100ChunkHKernel:
cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for k in cutlass.range_constexpr(BT // 16):
vtmem = v_tmem_base + k * 8
kdesc = kdesc_base | ((k * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True)
_tcgen05.commit(vk_done_mbar + stage_id)
for k in cutlass.range_constexpr(BT // 16):
vtmem = v_tmem_base + k * 8
kdesc = kdesc_base | ((k * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True)
_tcgen05.commit(vk_done_mbar + stage_id)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
@@ -421,7 +416,7 @@ class Sm100ChunkHKernel:
).load()
)
for j in cutlass.range_constexpr(32):
for j in cutlass.range(32, vectorize=True):
h_f32[j] *= h_scale
_tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
@@ -435,8 +430,10 @@ class Sm100ChunkHKernel:
fence_before_tma_store()
if warp_id_ == 3:
h_src = sH if cutlass.const_expr(is_f32) else sH0
h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
simple_tma_copy(H_tma_atom, h_src, h_dst)
h_dst = H_tma.tma_tensor[
chunk_offset + chunk_id, head_id, None, None
]
simple_tma_copy(H_tma.atom, h_src, h_dst)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
@@ -490,7 +487,7 @@ class Sm100ChunkHKernel:
h_f32.store(
_tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
)
for j in cutlass.range_constexpr(32):
for j in cutlass.range(32, vectorize=True):
h_f32[j] *= h_scale
_tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
_tcgen05.wait_st()
@@ -501,8 +498,10 @@ class Sm100ChunkHKernel:
cute.arch.barrier(barrier_id=1, number_of_threads=128)
fence_before_tma_store()
if warp_id_ == 3:
h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
simple_tma_copy(H_tma_atom, sH, h_dst)
h_dst = H_tma.tma_tensor[
chunk_offset + chunk_id, head_id, None, None
]
simple_tma_copy(H_tma.atom, sH, h_dst)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
@@ -530,8 +529,8 @@ class Sm100ChunkHKernel:
cute.arch.barrier(barrier_id=1, number_of_threads=128)
if warp_id_ == 0:
ht_dst = tmaHT[seq_id, head_id, None, None]
simple_tma_copy(HT_tma_atom, sH0, ht_dst)
ht_dst = HT_tma.tma_tensor[seq_id, head_id, None, None]
simple_tma_copy(HT_tma.atom, sH0, ht_dst)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
if warp_id_ == 1:
@@ -551,7 +550,7 @@ class Sm100ChunkHKernel:
# ((BT, num_BT_tiles), V_dim)
gV_new_tiles = cute.logical_divide(
tmaV_new[None, head_id, None], (BT, None)
V_new_tma.tma_tensor[None, head_id, None], (BT, None)
)
# sV shape: [BT, (64, V_dim/64), num_stages]
@@ -653,7 +652,7 @@ class Sm100ChunkHKernel:
fence_before_tma_store()
if warp_id == 3:
gV = gV_new_tiles[(None, chunk_offset + chunk_id), None]
simple_tma_copy(V_new_tma_atom, sV_new, gV)
simple_tma_copy(V_new_tma.atom, sV_new, gV)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
+297 -231
View File
@@ -11,6 +11,8 @@ from quack.compile_utils import make_fake_tensor
from vllm.cute_utils import (
EVICT_FIRST,
_bf16x2_neg,
_bf16x2_sub,
_tcgen05,
cvt,
fence_before_tma_store,
@@ -47,7 +49,7 @@ class Sm100ChunkUWKernel:
# hard-code
self.BT = 64
self.num_warps = 2 + 4 + 4
self.num_warps = 4 + 4 + 4
@cute.jit
def _make_tma_args(
@@ -68,13 +70,12 @@ class Sm100ChunkUWKernel:
# we need to convert gmem layout to (T, H, (64, D/64)) for make_tiled_tma_atom()
# to emit a single 4D TMA. otherwise, it will emit (D/64)x 3D TMA.
atom, tma_tensor = cpasync.make_tiled_tma_atom(
return cpasync.make_tiled_tma_atom(
op,
cute.logical_divide(tensor, (None, None, 64)),
slayout,
cta_tiler=(self.BT, 1, dim),
)
return atom, tma_tensor, slayout
@cute.jit
def __call__(
@@ -95,33 +96,33 @@ class Sm100ChunkUWKernel:
tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
K_args = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s)
V_args = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s)
U_args = self._make_tma_args(U, self.V_dim, 1, tma_s2g)
W_args = self._make_tma_args(W, self.K_dim, 1, tma_s2g)
K_tma = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s)
V_tma = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s)
U_tma = self._make_tma_args(U, self.V_dim, 1, tma_s2g)
W_tma = self._make_tma_args(W, self.K_dim, 1, tma_s2g)
grid = (num_sms // self.Hv, self.Hv, 1)
block = (self.num_warps * 32, 1, 1)
self.kernel(
K_args,
V_args,
U_args,
W_args,
K_tma,
V_tma,
U_tma,
W_tma,
g,
beta,
g_cu,
cu_seqlens,
chunk_indices,
total_chunks,
).launch(grid=grid, block=block, stream=stream)
).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
@cute.kernel
def kernel(
self,
K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
U_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
K_tma: cpasync.TmaInfo,
V_tma: cpasync.TmaInfo,
U_tma: cpasync.TmaInfo,
W_tma: cpasync.TmaInfo,
g: cute.Tensor,
beta: cute.Tensor,
g_cu: cute.Tensor,
@@ -142,10 +143,10 @@ class Sm100ChunkUWKernel:
V_dim = self.V_dim
num_stages = self.num_stages
K_tma_atom, tmaK, sK_layout = K_args
V_tma_atom, tmaV, sV_layout = V_args
U_tma_atom, tmaU, sU_layout = U_args
W_tma_atom, tmaW, sW_layout = W_args
INV_BAR = 1
EPI_BAR = 2
SCAN_BAR = 3
TMEM_ALLOC_BAR = 4
def allocate_tensor(smem, dtype, layout):
return smem.allocate_tensor(
@@ -153,22 +154,20 @@ class Sm100ChunkUWKernel:
)
smem = cutlass.utils.SmemAllocator()
sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
sU = allocate_tensor(smem, BFloat16, sU_layout)[None, 0, None, 0]
sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, 0]
sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
sU = allocate_tensor(smem, BFloat16, U_tma.smem_layout)[None, 0, None, 0]
sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, 0]
swizzle_128B = cute.make_swizzle(3, 4, 3)
sA_layout = cute.make_layout((BT, (64, 1)), stride=(64, (1, BT * 64)))
sA_layout = cute.make_composed_layout(swizzle_128B, 0, sA_layout)
sA = allocate_tensor(smem, BFloat16, sA_layout)
sAi = allocate_tensor(smem, BFloat16, sA_layout)
sA_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16)
sAi_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16)
s_beta = smem.allocate_array(Float32, BT)
s_g_cu_exp = smem.allocate_array(Float32, BT)
s_g_cu = smem.allocate_array(Float32, BT)
s_beta = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
s_g_cu = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
s_beta_g = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
tma_mbar = smem.allocate_array(Int64, num_stages)
prep_mbar = smem.allocate_array(Int64, num_stages)
mma_kkt_mbar = smem.allocate_array(Int64, num_stages)
inv_mbar = smem.allocate_array(Int64, num_stages)
mma_u_mbar = smem.allocate_array(Int64, num_stages)
@@ -179,7 +178,7 @@ class Sm100ChunkUWKernel:
kkt_tmem = 0
U_tmem_base = kkt_tmem + BT
Ab_tmem_base = U_tmem_base + V_dim * num_stages
assert Ab_tmem_base + (BT // 2) * num_stages <= 512
assert Ab_tmem_base + (BT // 2) <= 512
# prepare ldmatrix/stmatrix ops
ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4)
@@ -193,6 +192,7 @@ class Sm100ChunkUWKernel:
with cute.arch.elect_one():
for i in cutlass.range_constexpr(num_stages):
cute.arch.mbarrier_init(tma_mbar + i, 1)
cute.arch.mbarrier_init(prep_mbar + i, 64)
cute.arch.mbarrier_init(mma_kkt_mbar + i, 1)
cute.arch.mbarrier_init(inv_mbar + i, 128)
cute.arch.mbarrier_init(mma_u_mbar + i, 1)
@@ -200,14 +200,14 @@ class Sm100ChunkUWKernel:
cute.arch.mbarrier_init(epi_mbar + i, 128)
cute.arch.mbarrier_init_fence()
elif warp_id == 1:
cpasync.prefetch_descriptor(K_tma_atom)
cpasync.prefetch_descriptor(V_tma_atom)
cpasync.prefetch_descriptor(U_tma_atom)
cpasync.prefetch_descriptor(W_tma_atom)
cpasync.prefetch_descriptor(K_tma.atom)
cpasync.prefetch_descriptor(V_tma.atom)
cpasync.prefetch_descriptor(U_tma.atom)
cpasync.prefetch_descriptor(W_tma.atom)
cute.arch.sync_threads()
num_global_chunks = total_chunks[0]
if warp_id == 9:
if warp_id == 11:
# TMA warp
stage_id = 0
parity = 1
@@ -221,12 +221,14 @@ class Sm100ChunkUWKernel:
# domain_offset() to shift the pointer first.
mbar = tma_mbar + stage_id
gK = cute.local_tile(
cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
cute.domain_offset(
(bos, 0), K_tma.tma_tensor[None, k_head_id, None]
),
tiler=(BT, K_dim),
coord=(chunk_id, 0),
)
gV = cute.local_tile(
cute.domain_offset((bos, 0), tmaV[None, head_id, None]),
cute.domain_offset((bos, 0), V_tma.tma_tensor[None, head_id, None]),
tiler=(BT, V_dim),
coord=(chunk_id, 0),
)
@@ -237,18 +239,19 @@ class Sm100ChunkUWKernel:
with cute.arch.elect_one():
STAGE_SIZE = BT * (K_dim + V_dim) * 2
cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar)
simple_tma_copy(
V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
parity ^= 1
elif warp_id == 8:
elif warp_id == 10:
# MMA warp
_tcgen05.alloc(taddr)
cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160)
stage_id = 0
parity = 0
@@ -263,8 +266,8 @@ class Sm100ChunkUWKernel:
for global_chunk_id in range(bid, num_global_chunks, grid_x):
U_tmem = U_tmem_base + V_dim * stage_id
W_tmem = U_tmem | (16 << 16)
Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id
Abg_tmem = Ab_tmem | (16 << 16)
Ab_tmem = Ab_tmem_base
Abg_tmem = Ab_tmem_base | (16 << 16)
##### KKT MMA: KKT = K @ K.T #####
kaddr = sK[None, None, stage_id].iterator.toint()
@@ -276,18 +279,12 @@ class Sm100ChunkUWKernel:
cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_f16(
kkt_tmem,
kdesc,
kdesc,
kkt_idesc,
(i > 0) or (j > 0),
)
_tcgen05.commit(mma_kkt_mbar + stage_id)
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
enable_d = (i > 0) or (j > 0)
_tcgen05.mma_f16(kkt_tmem, kdesc, kdesc, kkt_idesc, enable_d)
_tcgen05.commit(mma_kkt_mbar + stage_id)
##### U/W MMA: U = Ab @ V, W = Abg @ K #####
vaddr = sV[None, None, stage_id].iterator.toint()
@@ -299,20 +296,15 @@ class Sm100ChunkUWKernel:
cute.arch.mbarrier_wait(inv_mbar + stage_id, parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(BT // 16):
_tcgen05.mma_ts_f16(
W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0
)
kdesc += (16 * 128) >> 4
_tcgen05.commit(mma_w_mbar + stage_id)
for i in cutlass.range_constexpr(BT // 16):
_tcgen05.mma_ts_f16(W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0)
kdesc += (16 * 128) >> 4
_tcgen05.commit(mma_w_mbar + stage_id)
for i in cutlass.range_constexpr(BT // 16):
_tcgen05.mma_ts_f16(
U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0
)
vdesc += (16 * 128) >> 4
_tcgen05.commit(mma_u_mbar + stage_id)
for i in cutlass.range_constexpr(BT // 16):
_tcgen05.mma_ts_f16(U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0)
vdesc += (16 * 128) >> 4
_tcgen05.commit(mma_u_mbar + stage_id)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
@@ -321,24 +313,124 @@ class Sm100ChunkUWKernel:
cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
_tcgen05.dealloc()
elif warp_id >= 8:
# dedicated prep warps for beta and gate, consumed by INV warps
stage_id = 0
parity = 0
tid_ = tid % 128
warp_id_ = warp_id % 4
for global_chunk_id in range(bid, num_global_chunks, grid_x):
seq_id = chunk_indices[global_chunk_id, 0]
chunk_id = chunk_indices[global_chunk_id, 1]
bos = cu_seqlens[seq_id]
eos = cu_seqlens[seq_id + 1]
off_t = bos + chunk_id * BT
t = off_t + tid_
in_bounds = t < eos
beta_val = beta[t, head_id] if in_bounds else Float32(0.0)
g_val = g[t, head_id] if in_bounds else Float32(0.0)
# warp-local prefix scan
for i in cutlass.range_constexpr(5):
offset = cutlass.const_expr(1 << i)
lower = cute.arch.shuffle_sync_up(g_val, offset, mask_and_clamp=0)
if lane_id >= offset:
g_val += lower
# Delay the stage-reuse wait until just before touching smem:
# global loads and the warp-local scan do not use staged buffers.
if warp_id_ == 0:
cute.arch.mbarrier_wait(inv_mbar + stage_id, parity ^ 1)
cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
# Store beta and the per-warp scan totals for the cross-warp fixup.
s_beta[tid_, stage_id] = beta_val
if lane_id == 31:
s_g_cu[warp_id_, stage_id] = g_val
cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
# Add the sum from the lower prep warp.
if warp_id_ == 1:
g_val += s_g_cu[0, stage_id]
cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
if in_bounds:
g_cu[t, head_id] = g_val
s_g_cu[tid_, stage_id] = g_val
s_beta_g[tid_, stage_id] = beta_val * cute.math.exp(g_val)
cute.arch.mbarrier_arrive(prep_mbar + stage_id)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
parity ^= 1
elif warp_id >= 4:
# inv warps
tid_ = tid % 128
warp_id_ = warp_id % 4
def store_ab_abg(
Ai_f32, s_beta, s_beta_g, warp_id_, lane_id, tile_col, Ab_tmem_base
):
# compute Ab and Abg from Ai, then store to tmem
beta_col = cute.make_rmem_tensor((2, 2), Float32)
beta_g_col = cute.make_rmem_tensor((2, 2), Float32)
for i in cutlass.range_constexpr(2):
base = tile_col * 16 + i * 8 + (lane_id % 4) * 2
for j in cutlass.range_constexpr(2):
beta_col[j, i] = s_beta[base + j]
beta_g_col[j, i] = s_beta_g[base + j]
# without conversion to TensorSSA, cutlass.range(vectorize=True) fails
beta_col = beta_col.load()
beta_g_col = beta_g_col.load()
Ab_f32 = cute.make_rmem_tensor(8, Float32)
Abg_f32 = cute.make_rmem_tensor(8, Float32)
for i in cutlass.range(8, vectorize=True):
scale_idx = (i // 4) * 2 + (i % 2)
Ab_f32[i] = Ai_f32[i] * beta_col[scale_idx]
Abg_f32[i] = Ai_f32[i] * beta_g_col[scale_idx]
Ab = Ab_f32.load().to(BFloat16)
Abg = Abg_f32.load().to(BFloat16)
Ab_tmem = Ab_tmem_base + tile_col * 8
_tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab)
_tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg)
# clear the Ab/Abg tmem buffer once before mainloop
# this is to keep the upper triangular tiles zeros
cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160)
zeros_ = cute.make_rmem_tensor(BT // 2, Float32)
zeros_.fill(0.0)
_tcgen05.st(warp_id_ * 32, Ab_tmem_base, "32x32b", BT // 2, zeros_)
stage_id = 0
parity = 0
# view into (16,16) sub-tiles, then ldmatrix layout
sA_ldsm = cute.logical_divide(sA, (16, cute.make_layout((8, 2))))
sAi_ldsm = cute.logical_divide(sAi, (16, cute.make_layout((8, 2))))
sA_ldsm = sA_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
sAi_ldsm = sAi_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
# for sA, we can avoid bank conflict without using swizzling because
# this is only used as tmp buffer between rmem<->smem, no gmem interactions.
# to do so, we can put (8,8) tile contiguous in memory. logically, we are
# partitioning (64,64) tile into 4x (16,16) tiles. to make indexing easier
# later, we view (16,16) tile as (32,8) tile here.
# (4,4) is (row_tile, col_tile).
sA_layout = cute.make_layout(((8, 32), (4, 4)))
sA = cute.make_tensor(sA_ptr, sA_layout)
sAi = cute.make_tensor(sAi_ptr, sA_layout)
# pre-compute ldmatrix addresses
sA_ldsm = sA[(None, lane_id), None]
sAi_ldsm = sAi[(None, lane_id), None]
# init Ai smem buffer with zeros (only the first 48 rows)
for i in cutlass.range_constexpr((BT // 4 * 3) * BT // 128):
idx = i * 128 + tid_
sAi[idx // BT, idx % BT] = BFloat16(0.0)
zeros_bf16 = cute.make_rmem_tensor(8, BFloat16)
zeros_bf16.fill(0.0)
for i in cutlass.range_constexpr(3):
cute.copy(stsm_atom, zeros_bf16, sAi_ldsm[None, (i, warp_id_)])
# indices for ldmatrix layout later
row_indices = cute.make_rmem_tensor((1, 2, 1), Int32)
@@ -360,61 +452,31 @@ class Sm100ChunkUWKernel:
eos = cu_seqlens[seq_id + 1]
off_t = bos + chunk_id * BT
t = off_t + tid_
##### Phase 1: load g and beta #####
if tid_ < BT:
in_bounds = t < eos
beta_val = beta[t, head_id] if in_bounds else Float32(0.0)
g_val = g[t, head_id] if in_bounds else Float32(0.0)
s_beta[tid_] = beta_val
# compute cumsum(g)
# parallel scan within a warp
for i in cutlass.range_constexpr(5):
offset = cutlass.const_expr(1 << i)
lower = cute.arch.shuffle_sync_up(
g_val, offset, mask_and_clamp=0
)
if lane_id >= offset:
g_val += lower
# store warp sum
if lane_id == 31:
s_g_cu[warp_id_] = g_val
cute.arch.barrier(barrier_id=3, number_of_threads=BT)
# add warp sum from lower warps
for i in cutlass.range_constexpr(1, BT // 32):
if warp_id_ >= i:
g_val += s_g_cu[i - 1]
cute.arch.barrier(barrier_id=3, number_of_threads=BT)
# store g_cu to gmem for H and O kernels
if in_bounds:
g_cu[t, head_id] = g_val
# store g and g_cu to smem for later
s_g_cu[tid_] = g_val
s_g_cu_exp[tid_] = cute.math.exp(g_val) if in_bounds else 0.0
##### Phase 2: A = strictLower(beta * kkt * Gamma) #####
##### Phase 1: A = strictLower(beta * kkt * Gamma) #####
# Ab/Abg share one tmem slot across stages. The MMA warp commits
# tcgen05 groups in program order: KKT_i, W_i, U_i, KKT_{i+1}.
# Waiting for KKT_i means the previous W/U commits have completed,
# so the INV warps can safely overwrite Ab/Abg for this iteration.
# Wait for prep warps to publish beta/gate and for KKT MMA.
if warp_id_ == 0:
cute.arch.mbarrier_wait(prep_mbar + stage_id, parity)
cute.arch.mbarrier_wait(mma_kkt_mbar + stage_id, parity)
cute.arch.barrier(barrier_id=1, number_of_threads=128)
cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
_tcgen05.fence_after_thread_sync()
# tmem 16x256b layout / ldmatrix layout
beta_row = cute.make_rmem_tensor(2, Float32)
g_cu_row = cute.make_rmem_tensor(2, Float32)
for i in cutlass.range_constexpr(2):
idx = warp_id_ * 16 + i * 8 + (lane_id // 4)
beta_row[i] = s_beta[idx, stage_id]
g_cu_row[i] = s_g_cu[idx, stage_id]
# mode0 is 8 rows together
# mode1 is top and bottom 8 rows
# mode2 is groups of 16 rows
row_coord = (lane_id // 4, None, warp_id_)
s_beta_view = cute.make_tensor(s_beta, (8, 2, 4))
beta_row = s_beta_view[row_coord].load().reshape((1, 2, 1))
s_g_cu_view = cute.make_tensor(s_g_cu, (8, 2, 4))
g_cu_row = s_g_cu_view[row_coord].load().reshape((1, 2, 1))
beta_row = beta_row.load().reshape((1, 2, 1))
g_cu_row = g_cu_row.load().reshape((1, 2, 1))
# mode0 is 2 consecutive elems
# mode1 is top and bottom 8 rows
@@ -429,7 +491,9 @@ class Sm100ChunkUWKernel:
# mode2 is top and bottom 8 rows
# mode3 is next 16 columns
col_coord = (None, lane_id % 4, None, i)
s_g_cu_view = cute.make_tensor(s_g_cu, (2, 4, 2, BT // 16))
s_g_cu_view = cute.make_tensor(
s_g_cu[None, stage_id].iterator, (2, 4, 2, BT // 16)
)
g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
Gamma = cute.math.exp(g_cu_row - g_cu_col, fastmath=True)
@@ -443,29 +507,23 @@ class Sm100ChunkUWKernel:
# pack to BF16
# CuteDSL doesn't generate cvt.bf16x2.f32 here for some reasons
packed = cute.make_rmem_tensor(4, Uint32)
packed[0] = cvt.fp32x2_to_bf16x2(
A_masked[0, 0, 0], A_masked[1, 0, 0]
)
packed[1] = cvt.fp32x2_to_bf16x2(
A_masked[0, 1, 0], A_masked[1, 1, 0]
)
packed[2] = cvt.fp32x2_to_bf16x2(
A_masked[0, 0, 1], A_masked[1, 0, 1]
)
packed[3] = cvt.fp32x2_to_bf16x2(
A_masked[0, 1, 1], A_masked[1, 1, 1]
)
for j in cutlass.range_constexpr(4):
packed[j] = cvt.fp32x2_to_bf16x2(
A_masked[j * 2], A_masked[j * 2 + 1]
)
# store to smem
cute.copy(
stsm_atom,
cute.recast_tensor(packed, BFloat16),
sA_ldsm[warp_id_, None, i],
sA_ldsm[None, (warp_id_, i)],
)
cute.arch.barrier(barrier_id=1, number_of_threads=128)
# use sync warp instead of bar.sync because for block-diagonal inverse,
# each warp reads its own private smem memory.
cute.arch.sync_warp()
##### Phase 3: matrix inverse #####
##### Phase 2: matrix inverse #####
# we use Newton-Schulz iterations to compute the inverse
# of the four 16x16 diagonal blocks.
# Ai_new = 2 Ai - Ai @ M @ Ai
@@ -477,15 +535,6 @@ class Sm100ChunkUWKernel:
zeros_f32 = cute.make_rmem_tensor(4, Float32)
zeros_f32.fill(0.0)
def set_diagonal(A: cute.Tensor, lane_id: Int32):
"Set the diagonal to 1s"
if lane_id % 9 == 0:
A[0] = (A[0] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
A[3] = (A[3] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
elif lane_id % 9 == 4:
A[0] = (A[0] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
A[3] = (A[3] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
Ai_bf16 = cute.make_rmem_tensor(8, BFloat16)
mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16)
M_bf16 = cute.make_rmem_tensor(8, BFloat16)
@@ -496,44 +545,60 @@ class Sm100ChunkUWKernel:
mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2)
M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2)
# construct rmem-backed identity matrix
eye = cute.make_rmem_tensor(4, Uint32)
eye[0] = Uint32(lane_id % 9 == 0) * Uint32(0x00003F80) + Uint32(
lane_id % 9 == 4
) * Uint32(0x3F800000)
eye[1] = 0
eye[2] = 0
eye[3] = eye[0]
# initial guess: Ai = I-A
cute.copy(ldsm_atom, sA_ldsm[warp_id_, None, warp_id_], Ai_bf16)
cute.copy(ldsm_atom, sA_ldsm[None, (warp_id_, warp_id_)], Ai_bf16)
for i in cutlass.range_constexpr(4):
Ai[i] ^= Uint32(0x80008000) # negate A
set_diagonal(Ai, lane_id)
Ai[i] = _bf16x2_sub(eye[i], Ai[i])
# (4, 2)
Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4)
# M is holding -(I+A), stay constant throughout the iterations
cute.copy(ldsm_trans_atom, sA_ldsm[warp_id_, None, warp_id_], M_bf16)
set_diagonal(M, lane_id)
cute.copy(ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], M_bf16)
for i in cutlass.range_constexpr(4):
M[i] ^= Uint32(0x80008000)
M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i])
# 3 rounds of Newton-Schulz
for _ in cutlass.range_constexpr(3):
# First MMA: -AiM = Ai @ (-M)
cute.copy(stsm_atom, Ai_bf16, sA_ldsm[warp_id_, None, warp_id_])
cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)])
cute.arch.sync_warp()
acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32)
Ai_bf16.store(acc.load().to(BFloat16))
# Second MMA: Ai_new = 2Ai + (-AiM) @ Ai
for j in cutlass.range_constexpr(8):
for j in cutlass.range(8, vectorize=True):
Ai_f32[j] *= 2.0
cute.copy(
ldsm_trans_atom,
sA_ldsm[warp_id_, None, warp_id_],
sA_ldsm[None, (warp_id_, warp_id_)],
mma_B_bf16,
)
Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0])
Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1])
Ai_bf16.store(Ai_f32.load().to(BFloat16))
cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[warp_id_, None, warp_id_])
cute.arch.barrier(barrier_id=1, number_of_threads=128)
cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[None, (warp_id_, warp_id_)])
store_ab_abg(
Ai_f32,
s_beta[None, stage_id],
s_beta_g[None, stage_id],
warp_id_,
lane_id,
warp_id_,
Ab_tmem_base,
)
cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
# off-diagonal by 1
# given
@@ -544,14 +609,14 @@ class Sm100ChunkUWKernel:
# warp1: Ai10 = -Ai11 @ A10 @ Ai00
# warp2: Ai21 = -Ai22 @ A21 @ Ai11
# warp3: Ai32 = -Ai33 @ A32 @ Ai22
if warp_id_ > 0:
if warp_id_ >= 1:
neg_Ai = cute.make_rmem_tensor(4, Uint32)
for i in cutlass.range_constexpr(4):
neg_Ai[i] = Ai[i] ^ Uint32(0x80008000)
neg_Ai[i] = _bf16x2_neg(Ai[i])
cute.copy(
ldsm_trans_atom,
sA_ldsm[warp_id_, None, warp_id_ - 1],
sA_ldsm[None, (warp_id_, warp_id_ - 1)],
mma_B_bf16,
)
acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32)
@@ -560,44 +625,52 @@ class Sm100ChunkUWKernel:
cute.copy(
ldsm_trans_atom,
sAi_ldsm[warp_id_ - 1, None, warp_id_ - 1],
sAi_ldsm[None, (warp_id_ - 1, warp_id_ - 1)],
mma_B_bf16,
)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
Ai_bf16.store(acc.load().to(BFloat16))
store_ab_abg(
acc,
s_beta[None, stage_id],
s_beta_g[None, stage_id],
warp_id_,
lane_id,
warp_id_ - 1,
Ab_tmem_base,
)
cute.copy(
stsm_atom,
Ai_bf16,
sAi_ldsm[warp_id_, None, warp_id_ - 1],
sAi_ldsm[None, (warp_id_, warp_id_ - 1)],
)
cute.arch.barrier(barrier_id=1, number_of_threads=128)
cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
# off-diagonal by 2
# warp0: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10)
# warp1: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21)
if warp_id_ < 2:
# warp2: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10)
# warp3: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21)
if warp_id_ >= 2:
tile_col = warp_id_ - 2
cute.copy(
ldsm_atom,
sA_ldsm[warp_id_ + 2, None, warp_id_],
sA_ldsm[None, (warp_id_, tile_col)],
Ai_bf16,
)
cute.copy(
ldsm_trans_atom,
sAi_ldsm[warp_id_, None, warp_id_],
sAi_ldsm[None, (tile_col, tile_col)],
mma_B_bf16,
)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
cute.copy(
ldsm_atom,
sA_ldsm[warp_id_ + 2, None, warp_id_ + 1],
Ai_bf16,
ldsm_atom, sA_ldsm[None, (warp_id_, tile_col + 1)], Ai_bf16
)
cute.copy(
ldsm_trans_atom,
sAi_ldsm[warp_id_ + 1, None, warp_id_],
sAi_ldsm[None, (tile_col + 1, tile_col)],
mma_B_bf16,
)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
@@ -605,79 +678,68 @@ class Sm100ChunkUWKernel:
tmp = cute.make_rmem_tensor(8, BFloat16)
tmp.store(acc.load().to(BFloat16))
cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)])
cute.arch.sync_warp()
cute.copy(
ldsm_atom, sAi_ldsm[warp_id_ + 2, None, warp_id_ + 2], Ai_bf16
)
cute.copy(ldsm_atom, sAi_ldsm[None, (warp_id_, warp_id_)], Ai_bf16)
for i in cutlass.range_constexpr(4):
Ai[i] ^= Uint32(0x80008000)
Ai[i] = _bf16x2_neg(Ai[i])
cute.copy(
ldsm_trans_atom,
sAi_ldsm[warp_id_ + 2, None, warp_id_],
sAi_ldsm[None, (warp_id_, tile_col)],
mma_B_bf16,
)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
tmp.store(acc.load().to(BFloat16))
cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
cute.arch.barrier(barrier_id=1, number_of_threads=128)
cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)])
store_ab_abg(
acc,
s_beta[None, stage_id],
s_beta_g[None, stage_id],
warp_id_,
lane_id,
tile_col,
Ab_tmem_base,
)
cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
# off-diagonal by 3
# warp0: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)
if warp_id_ == 0:
cute.copy(ldsm_atom, sA_ldsm[3, None, 0], Ai_bf16)
cute.copy(ldsm_trans_atom, sAi_ldsm[0, None, 0], mma_B_bf16)
# warp3: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)
if warp_id_ == 3:
cute.copy(ldsm_atom, sA_ldsm[None, (3, 0)], Ai_bf16)
cute.copy(ldsm_trans_atom, sAi_ldsm[None, (0, 0)], mma_B_bf16)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
for i in cutlass.range_constexpr(1, 3):
cute.copy(ldsm_atom, sA_ldsm[3, None, i], Ai_bf16)
cute.copy(ldsm_trans_atom, sAi_ldsm[i, None, 0], mma_B_bf16)
cute.copy(ldsm_atom, sA_ldsm[None, (3, i)], Ai_bf16)
cute.copy(ldsm_trans_atom, sAi_ldsm[None, (i, 0)], mma_B_bf16)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
tmp = cute.make_rmem_tensor(8, BFloat16)
tmp.store(acc.load().to(BFloat16))
cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)])
cute.arch.sync_warp()
cute.copy(ldsm_atom, sAi_ldsm[3, None, 3], Ai_bf16)
cute.copy(ldsm_atom, sAi_ldsm[None, (3, 3)], Ai_bf16)
for i in cutlass.range_constexpr(4):
Ai[i] ^= Uint32(0x80008000)
cute.copy(ldsm_trans_atom, sAi_ldsm[3, None, 0], mma_B_bf16)
Ai[i] = _bf16x2_neg(Ai[i])
cute.copy(ldsm_trans_atom, sAi_ldsm[None, (3, 0)], mma_B_bf16)
acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
tmp.store(acc.load().to(BFloat16))
cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
##### Phase 4: compute Ab, Abg #####
if warp_id_ == 3:
cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity ^ 1)
cute.arch.barrier(barrier_id=1, number_of_threads=128)
for i in cutlass.range_constexpr(BT // 16):
cute.copy(ldsm_atom, sAi_ldsm[warp_id_, None, i], Ai_bf16)
col_coord = (None, lane_id % 4, None, i)
s_beta_view = cute.make_tensor(s_beta, (2, 4, 2, BT // 16))
beta_col = s_beta_view[col_coord].load().reshape((2, 1, 2))
s_g_cu_view = cute.make_tensor(s_g_cu_exp, (2, 4, 2, BT // 16))
g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
Ai_f32 = cvt.bf16x2_to_fp32x2(Ai).load().reshape((2, 2, 2))
Ab_f32 = Ai_f32 * beta_col
Ab = Ab_f32.to(BFloat16)
Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id + i * 8
_tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab)
Abg_f32 = Ab_f32 * g_cu_col
Abg = Abg_f32.to(BFloat16)
_tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg)
cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)])
store_ab_abg(
acc,
s_beta[None, stage_id],
s_beta_g[None, stage_id],
warp_id_,
lane_id,
0,
Ab_tmem_base,
)
_tcgen05.wait_st()
_tcgen05.fence_before_thread_sync()
cute.arch.mbarrier_arrive(inv_mbar + stage_id)
@@ -686,14 +748,18 @@ class Sm100ChunkUWKernel:
if stage_id == 0:
parity ^= 1
elif warp_id < 4:
else:
# epi warps
stage_id = 0
parity = 0
# ((BT, num_global_chunks), V_dim)
gU_tiles = cute.logical_divide(tmaU[None, head_id, None], (BT, None))
gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
gU_tiles = cute.logical_divide(
U_tma.tma_tensor[None, head_id, None], (BT, None)
)
gW_tiles = cute.logical_divide(
W_tma.tma_tensor[None, head_id, None], (BT, None)
)
# sW shape: [BT, (64, K_dim/64)]
# sW_view shape: [(8, 2), (4, K_dim/64)]
@@ -719,7 +785,7 @@ class Sm100ChunkUWKernel:
elif warp_id == 1:
with cute.arch.elect_one():
cute.arch.cp_async_bulk_wait_group(0, read=True)
cute.arch.barrier(barrier_id=2, number_of_threads=128)
cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
_tcgen05.fence_after_thread_sync()
w_f32 = _tcgen05.ld(warp_id * 32 + 16, U_tmem, "16x256b", K_dim // 8)
@@ -729,16 +795,16 @@ class Sm100ChunkUWKernel:
cute.copy(stsm_atom, w_bf16, sW_view)
# wait for U MMA + issue W TMA store
cute.arch.barrier(barrier_id=2, number_of_threads=128)
cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
fence_before_tma_store()
if warp_id == 0:
cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
elif warp_id == 1:
# don't need to commit
simple_tma_copy(
W_tma_atom, sW, gW_tiles[(None, global_chunk_id), None]
W_tma.atom, sW, gW_tiles[(None, global_chunk_id), None]
)
cute.arch.barrier(barrier_id=2, number_of_threads=128)
cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
_tcgen05.fence_after_thread_sync()
u_f32 = _tcgen05.ld(warp_id * 32, U_tmem, "16x256b", V_dim // 8)
@@ -749,11 +815,11 @@ class Sm100ChunkUWKernel:
u_bf16.store(u_f32.to(BFloat16))
cute.copy(stsm_atom, u_bf16, sU_view)
cute.arch.barrier(barrier_id=2, number_of_threads=128)
cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
fence_before_tma_store()
if warp_id == 1:
simple_tma_copy(
U_tma_atom, sU, gU_tiles[(None, global_chunk_id), None]
U_tma.atom, sU, gU_tiles[(None, global_chunk_id), None]
)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py
+82 -80
View File
@@ -61,13 +61,12 @@ class Sm100ChunkOKernel:
stride=(64, 0, (1, self.BT * 64), self.BT * dim),
)
slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
atom, tma_tensor = cpasync.make_tiled_tma_atom(
return cpasync.make_tiled_tma_atom(
op,
cute.logical_divide(tensor, (None, None, 64)),
slayout,
cta_tiler=(self.BT, 1, dim),
)
return atom, tma_tensor, slayout
@cute.jit
def _make_h_tma_args(
@@ -79,17 +78,22 @@ class Sm100ChunkOKernel:
num_elems = 128 // (tensor.element_type.width // 8)
swizzle_128B = cute.make_swizzle(3, 4, 3)
slayout = cute.make_layout(
(1, self.V_dim, (num_elems, self.K_dim // num_elems), stages),
stride=(0, num_elems, (1, self.V_dim * num_elems), self.V_dim * self.K_dim),
(1, 1, self.V_dim, (num_elems, self.K_dim // num_elems), stages),
stride=(
0,
0,
num_elems,
(1, self.V_dim * num_elems),
self.V_dim * self.K_dim,
),
)
slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
atom, tma_tensor = cpasync.make_tiled_tma_atom(
return cpasync.make_tiled_tma_atom(
op,
cute.logical_divide(tensor, (None, None, num_elems)),
cute.logical_divide(tensor, (None, None, None, num_elems)),
slayout,
cta_tiler=(1, self.V_dim, self.K_dim),
cta_tiler=(1, 1, self.V_dim, self.K_dim),
)
return atom, tma_tensor, slayout
@cute.jit
def __call__(
@@ -111,35 +115,35 @@ class Sm100ChunkOKernel:
block = (self.num_warps * 32, 1, 1)
tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
Q_args = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages)
K_args = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages)
V_args = self._make_bf16_tma_args(
Q_tma = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages)
K_tma = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages)
V_tma = self._make_bf16_tma_args(
v_new_chunks, self.V_dim, tma_g2s, self.num_stages
)
H_args = self._make_h_tma_args(h, tma_g2s, self.num_stages)
O_args = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1)
H_tma = self._make_h_tma_args(h, tma_g2s, self.num_stages)
O_tma = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1)
self.kernel(
Q_args,
K_args,
V_args,
H_args,
O_args,
Q_tma,
K_tma,
V_tma,
H_tma,
O_tma,
g_cu,
o,
cu_seqlens,
chunk_indices,
total_chunks,
scale,
).launch(grid=grid, block=block, stream=stream)
).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
@cute.kernel
def kernel(
self,
Q_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
O_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
Q_tma: cpasync.TmaInfo,
K_tma: cpasync.TmaInfo,
V_tma: cpasync.TmaInfo,
H_tma: cpasync.TmaInfo,
O_tma: cpasync.TmaInfo,
g_cu: cute.Tensor,
o: cute.Tensor,
cu_seqlens: cute.Tensor,
@@ -162,23 +166,17 @@ class Sm100ChunkOKernel:
k_head_id = v_head_id // heads_per_qk
num_global_chunks = total_chunks[0]
Q_tma_atom, tmaQ, sQ_layout = Q_args
K_tma_atom, tmaK, sK_layout = K_args
V_tma_atom, tmaV, sV_layout = V_args
H_tma_atom, tmaH, sH_layout = H_args
O_tma_atom, tmaO, sO_layout = O_args
def allocate_tensor(smem, dtype, layout):
return smem.allocate_tensor(
dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
)
smem = cutlass.utils.SmemAllocator()
sQ = allocate_tensor(smem, BFloat16, sQ_layout)[None, 0, None, None]
sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
sH = allocate_tensor(smem, BFloat16, sH_layout)[0, None, None, None]
sO = allocate_tensor(smem, BFloat16, sO_layout)[None, 0, None, 0]
sQ = allocate_tensor(smem, BFloat16, Q_tma.smem_layout)[None, 0, None, None]
sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None, None]
sO = allocate_tensor(smem, BFloat16, O_tma.smem_layout)[None, 0, None, 0]
s_g_cu = smem.allocate_array(Float32, BT)
qk_full_mbar = smem.allocate_array(Int64, num_stages)
@@ -207,10 +205,10 @@ class Sm100ChunkOKernel:
cute.arch.mbarrier_init(epi_mbar, 128)
cute.arch.mbarrier_init_fence()
elif warp_id == 9:
cpasync.prefetch_descriptor(Q_tma_atom)
cpasync.prefetch_descriptor(K_tma_atom)
cpasync.prefetch_descriptor(V_tma_atom)
cpasync.prefetch_descriptor(H_tma_atom)
cpasync.prefetch_descriptor(Q_tma.atom)
cpasync.prefetch_descriptor(K_tma.atom)
cpasync.prefetch_descriptor(V_tma.atom)
cpasync.prefetch_descriptor(H_tma.atom)
cute.arch.sync_threads()
if warp_id == 9:
@@ -225,12 +223,16 @@ class Sm100ChunkOKernel:
# copy Q and K
q_tile = cute.local_tile(
cute.domain_offset((bos, 0), tmaQ[None, k_head_id, None]),
cute.domain_offset(
(bos, 0), Q_tma.tma_tensor[None, k_head_id, None]
),
tiler=(BT, K_dim),
coord=(chunk_id, 0),
)
k_tile = cute.local_tile(
cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
cute.domain_offset(
(bos, 0), K_tma.tma_tensor[None, k_head_id, None]
),
tiler=(BT, K_dim),
coord=(chunk_id, 0),
)
@@ -241,13 +243,13 @@ class Sm100ChunkOKernel:
with cute.arch.elect_one():
STAGE_SIZE = BT * (K_dim + K_dim) * 2
cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
simple_tma_copy(Q_tma_atom, q_tile, sQ[None, None, stage_id], mbar)
simple_tma_copy(K_tma_atom, k_tile, sK[None, None, stage_id], mbar)
simple_tma_copy(Q_tma.atom, q_tile, sQ[None, None, stage_id], mbar)
simple_tma_copy(K_tma.atom, k_tile, sK[None, None, stage_id], mbar)
# copy H and V
gH = tmaH[global_chunk_id * self.Hv + v_head_id, None, None]
gH = H_tma.tma_tensor[global_chunk_id, v_head_id, None, None]
gV = cute.local_tile(
tmaV[None, v_head_id, None],
V_tma.tma_tensor[None, v_head_id, None],
tiler=(BT, V_dim),
coord=(global_chunk_id, 0),
)
@@ -262,10 +264,10 @@ class Sm100ChunkOKernel:
mbar, H_STAGE_SIZE + V_STAGE_SIZE
)
simple_tma_copy(
H_tma_atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
H_tma.atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
)
simple_tma_copy(
V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
)
stage_id = (stage_id + 1) % num_stages
@@ -302,40 +304,37 @@ class Sm100ChunkOKernel:
cute.arch.mbarrier_wait(qk_full_mbar + stage_id, tma_parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(K_dim // BT):
for j in cutlass.range_constexpr(BT // 16):
qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_f16(
qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0)
)
_tcgen05.commit(qk_mbar)
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
_tcgen05.mma_f16(
qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0)
)
_tcgen05.commit(qk_mbar)
##### 2nd MMA: Q @ H.T #####
cute.arch.mbarrier_wait(hv_full_mbar + stage_id, tma_parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(K_dim // BT):
for j in cutlass.range_constexpr(BT // 16):
qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4)
_tcgen05.mma_f16(
qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0)
)
_tcgen05.commit(qk_empty_mbar + stage_id)
for i in cutlass.range_constexpr(K_dim // 64):
for j in cutlass.range_constexpr(64 // 16):
qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4)
_tcgen05.mma_f16(
qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0)
)
_tcgen05.commit(qk_empty_mbar + stage_id)
##### 3rd MMA: P @ V #####
# stalled by mask(QK)
cute.arch.mbarrier_wait(mask_mbar, mask_parity)
_tcgen05.fence_after_thread_sync()
with cute.arch.elect_one():
for i in cutlass.range_constexpr(BT // 16):
vdesc = vdesc_base | ((i * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(
out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0
)
_tcgen05.commit(pv_mma_mbar + stage_id)
for i in cutlass.range_constexpr(BT // 16):
vdesc = vdesc_base | ((i * 16 * 128) >> 4)
_tcgen05.mma_ts_f16(
out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0
)
_tcgen05.commit(pv_mma_mbar + stage_id)
stage_id = (stage_id + 1) % num_stages
if stage_id == 0:
@@ -499,11 +498,13 @@ class Sm100ChunkOKernel:
fence_before_tma_store()
if warp_id == 3:
gO = cute.local_tile(
cute.domain_offset((bos, 0), tmaO[None, v_head_id, None]),
cute.domain_offset(
(bos, 0), O_tma.tma_tensor[None, v_head_id, None]
),
tiler=(BT, V_dim),
coord=(chunk_id, 0),
)
simple_tma_copy(O_tma_atom, sO, gO)
simple_tma_copy(O_tma.atom, sO, gO)
with cute.arch.elect_one():
cute.arch.cp_async_bulk_commit_group()
@@ -559,13 +560,14 @@ class Sm100ChunkOKernel:
total_t = cute.sym_int()
pad_t = cute.sym_int()
total_chunks_n = cute.sym_int()
h_outer_n = cute.sym_int()
cu_entries = cute.sym_int()
q = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
k = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
v_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
h_flat = make_fake_tensor(BFloat16, (h_outer_n, V_dim, K_dim), divisibility=16)
h = make_fake_tensor(
BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
)
g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
o = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
@@ -586,7 +588,7 @@ class Sm100ChunkOKernel:
q,
k,
v_new,
h_flat,
h,
g_cu,
o,
cu_seqlens,
@@ -602,7 +604,7 @@ class Sm100ChunkOKernel:
def o_cutedsl(
q: torch.Tensor,
k: torch.Tensor,
v_new_chunks: torch.Tensor,
v_new: torch.Tensor,
h: torch.Tensor,
g_cu: torch.Tensor,
o: torch.Tensor,
@@ -618,8 +620,8 @@ def o_cutedsl(
Sm100ChunkOKernel.compile(H, Hv, K_dim, V_dim)(
q,
k,
v_new_chunks.view(-1, Hv, V_dim),
h.view(-1, V_dim, K_dim),
v_new,
h,
g_cu,
o,
cu_seqlens,