Skip to content

[None][fix] Double-buffer update_ready in the KDA decomposed chain kernel #22179

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ruocheng-nv wants to merge 1 commit into
base: main
Choose a base branch
from
+109 −12
Open

[None][fix] Double-buffer update_ready in the KDA decomposed chain kernel #22179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 96 additions & 0 deletions 3rdparty/vendor_patches/flashinfer-kda-bt16.patch
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,99 @@
diff --git a/kda_chunked_bt16.py b/kda_chunked_bt16.py
index 38903d429cbcfa05c3363a760f3c32902532e7cf..f90bf435e2d1bdf2a8d89c90a8bcb487e2fdcba0 100644
--- a/kda_chunked_bt16.py
+++ b/kda_chunked_bt16.py
@@ -8713,8 +8713,10 @@ def kernel_chain_dv2(
u_input_ready_mbar = cutlass.Array(
cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
)
+ # Two slots: warp 12 commits stateq_done before its wait, so one slot can
+ # complete twice before warp 12 polls it.
update_ready_mbar = cutlass.Array(
- cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
+ cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
)
shared_acc_ready_mbar = cutlass.Array(
cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
@@ -8823,7 +8825,8 @@ def kernel_chain_dv2(
prims.mbarrier_init(stateq_done_mbar.subview(stage), 1)
# CG1 fuses the right-half pack/scale while idle CG0 warps 4-7
# fuse the left half. Both groups must finish before delta.
- prims.mbarrier_init(update_ready_mbar, 8)
+ for stage in cutlass.range_constexpr(2):
+ prims.mbarrier_init(update_ready_mbar.subview(stage), 8)
for stage in cutlass.range_constexpr(K2_QSTATE_STAGE_COUNT):
prims.mbarrier_init(qstate_acc_ready_mbar.subview(stage), 1)
prims.mbarrier_init(final_state_stored_mbar, 8)
@@ -8921,7 +8924,6 @@ def kernel_chain_dv2(
elif warp_idx == ROLES.super_mma:
tmem_raw_addr = tmem_ptr_i32.load()
si_phase12 = cutlass.Int32(0)
- upd_phase12 = cutlass.Int32(0)
for chunk in cutlass.range(num_chunks, unroll=1):
raw_stage = chunk % K2_RAW_STAGE_COUNT
qstate_stage = chunk % K2_QSTATE_STAGE_COUNT
@@ -8958,7 +8960,7 @@ def kernel_chain_dv2(
True,
DV_HALF,
)
- upd_phase12 = update_ready_wait(update_ready_mbar, upd_phase12)
+ update_ready_wait(update_ready_mbar.subview(chunk % 2), (chunk // 2) % 2)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
xpack_col12 = tcgen05_shared_input_tmem_col_offset(acc_stage)
tcgen05_chain_issue_qkv_mma(
@@ -8976,7 +8978,6 @@ def kernel_chain_dv2(
tmem_raw_addr = tmem_ptr_i32.load()
si_l_phase = cutlass.Int32(0)
si_phase = cutlass.Int32(0)
- upd_phase = cutlass.Int32(0)
for chunk in cutlass.range(num_chunks, unroll=1):
raw_stage = chunk % K2_RAW_STAGE_COUNT
raw_phase = (chunk // K2_RAW_STAGE_COUNT) % 2
@@ -9018,7 +9019,7 @@ def kernel_chain_dv2(
)

xpack_col = tcgen05_shared_input_tmem_col_offset(acc_stage)
- upd_phase = update_ready_wait(update_ready_mbar, upd_phase)
+ update_ready_wait(update_ready_mbar.subview(chunk % 2), (chunk // 2) % 2)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
tcgen05_chain_issue_delta_half_mma(
w_stage_smem, tmem_raw_addr, xpack_col, input_dtype, 0
@@ -9153,7 +9154,7 @@ def kernel_chain_dv2(
input_dtype,
)
post_scale_lane = cute.arch.lane_idx()
- update_ready_arrive(update_ready_mbar)
+ update_ready_arrive(update_ready_mbar.subview(0))
if cutlass.const_expr(state_ckpt is not None):
# Peeled chunk 0's checkpoint (stride 1 only).
checkpoint_due = ckpt_next == cutlass.Int32(1)
@@ -9233,7 +9234,7 @@ def kernel_chain_dv2(
input_dtype,
)
post_scale_lane = cute.arch.lane_idx()
- update_ready_arrive(update_ready_mbar)
+ update_ready_arrive(update_ready_mbar.subview(chunk % 2))
if cutlass.const_expr(state_ckpt is not None):
# Same contract as the engine-side checkpoint; both M=64 TMEM
# halves wait their own k_restore parity slot (the same wait
@@ -9340,7 +9341,7 @@ def kernel_chain_dv2(
cutlass.Int32(0),
)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
- update_ready_arrive(update_ready_mbar)
+ update_ready_arrive(update_ready_mbar.subview(0))
for chunk in cutlass.range(1, num_chunks, 1, unroll=1):
prev = chunk - cutlass.Int32(1)
raw_stage = chunk % K2_RAW_STAGE_COUNT
@@ -9379,7 +9380,7 @@ def kernel_chain_dv2(
state_input_ready_l_mbar,
0,
)
- update_ready_arrive(update_ready_mbar)
+ update_ready_arrive(update_ready_mbar.subview(chunk % 2))


@cute.jit
diff --git a/kda_chunked_bt16_small_bh.py b/kda_chunked_bt16_small_bh.py
index e8bc9feb5ece3dbfe3156c90bb27e68fa5a4521a..65daa2b606343382b7b9e66a4997215b6aaa0e44 100644
--- a/kda_chunked_bt16_small_bh.py
Expand Down
4 changes: 2 additions & 2 deletions 3rdparty/vendor_sources.lock.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@ vendors:
- kda_chunked_bt16.py
- kda_chunked_bt16_small_bh.py
patch: 3rdparty/vendor_patches/flashinfer-kda-bt16.patch
patch_digest: sha256:42e8306de30fef448378066321977aeb7b42c3405dc883f9527e256590a58b46
digest: sha256-tree-v1:12eb5c011c3fef1831aeba02691df8e2f5bff323f9da9c58c0f31e80883a9f06
patch_digest: sha256:77729d1f7c3d29682d232e6128be8b1bfcb548cdcf54dd16995685d620a6eeea
digest: sha256-tree-v1:b5e8540bfd32251213a13d0aec255124a75606eb78cd1a9de5256a0062c9a9a9
flashinfer-prims-ts:
url: https://github.com/yuxianq/flashinfer.git
branch: trtllm-prims-ts-dev-20261107
Expand Down
21 changes: 11 additions & 10 deletions tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8713,8 +8713,10 @@ def kernel_chain_dv2(
u_input_ready_mbar = cutlass.Array(
cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
)
# Two slots: warp 12 commits stateq_done before its wait, so one slot can
# complete twice before warp 12 polls it.
update_ready_mbar = cutlass.Array(
cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8
cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
)
shared_acc_ready_mbar = cutlass.Array(
cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8
Expand Down Expand Up @@ -8823,7 +8825,8 @@ def kernel_chain_dv2(
prims.mbarrier_init(stateq_done_mbar.subview(stage), 1)
# CG1 fuses the right-half pack/scale while idle CG0 warps 4-7
# fuse the left half. Both groups must finish before delta.
prims.mbarrier_init(update_ready_mbar, 8)
for stage in cutlass.range_constexpr(2):
prims.mbarrier_init(update_ready_mbar.subview(stage), 8)
for stage in cutlass.range_constexpr(K2_QSTATE_STAGE_COUNT):
prims.mbarrier_init(qstate_acc_ready_mbar.subview(stage), 1)
prims.mbarrier_init(final_state_stored_mbar, 8)
Expand Down Expand Up @@ -8921,7 +8924,6 @@ def kernel_chain_dv2(
elif warp_idx == ROLES.super_mma:
tmem_raw_addr = tmem_ptr_i32.load()
si_phase12 = cutlass.Int32(0)
upd_phase12 = cutlass.Int32(0)
for chunk in cutlass.range(num_chunks, unroll=1):
raw_stage = chunk % K2_RAW_STAGE_COUNT
qstate_stage = chunk % K2_QSTATE_STAGE_COUNT
Expand Down Expand Up @@ -8958,7 +8960,7 @@ def kernel_chain_dv2(
True,
DV_HALF,
)
upd_phase12 = update_ready_wait(update_ready_mbar, upd_phase12)
update_ready_wait(update_ready_mbar.subview(chunk % 2), (chunk // 2) % 2)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
xpack_col12 = tcgen05_shared_input_tmem_col_offset(acc_stage)
tcgen05_chain_issue_qkv_mma(
Expand All @@ -8976,7 +8978,6 @@ def kernel_chain_dv2(
tmem_raw_addr = tmem_ptr_i32.load()
si_l_phase = cutlass.Int32(0)
si_phase = cutlass.Int32(0)
upd_phase = cutlass.Int32(0)
for chunk in cutlass.range(num_chunks, unroll=1):
raw_stage = chunk % K2_RAW_STAGE_COUNT
raw_phase = (chunk // K2_RAW_STAGE_COUNT) % 2
Expand Down Expand Up @@ -9018,7 +9019,7 @@ def kernel_chain_dv2(
)

xpack_col = tcgen05_shared_input_tmem_col_offset(acc_stage)
upd_phase = update_ready_wait(update_ready_mbar, upd_phase)
update_ready_wait(update_ready_mbar.subview(chunk % 2), (chunk // 2) % 2)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
tcgen05_chain_issue_delta_half_mma(
w_stage_smem, tmem_raw_addr, xpack_col, input_dtype, 0
Expand Down Expand Up @@ -9153,7 +9154,7 @@ def kernel_chain_dv2(
input_dtype,
)
post_scale_lane = cute.arch.lane_idx()
update_ready_arrive(update_ready_mbar)
update_ready_arrive(update_ready_mbar.subview(0))
if cutlass.const_expr(state_ckpt is not None):
# Peeled chunk 0's checkpoint (stride 1 only).
checkpoint_due = ckpt_next == cutlass.Int32(1)
Expand Down Expand Up @@ -9233,7 +9234,7 @@ def kernel_chain_dv2(
input_dtype,
)
post_scale_lane = cute.arch.lane_idx()
update_ready_arrive(update_ready_mbar)
update_ready_arrive(update_ready_mbar.subview(chunk % 2))
if cutlass.const_expr(state_ckpt is not None):
# Same contract as the engine-side checkpoint; both M=64 TMEM
# halves wait their own k_restore parity slot (the same wait
Expand Down Expand Up @@ -9340,7 +9341,7 @@ def kernel_chain_dv2(
cutlass.Int32(0),
)
prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC)
update_ready_arrive(update_ready_mbar)
update_ready_arrive(update_ready_mbar.subview(0))
for chunk in cutlass.range(1, num_chunks, 1, unroll=1):
prev = chunk - cutlass.Int32(1)
raw_stage = chunk % K2_RAW_STAGE_COUNT
Expand Down Expand Up @@ -9379,7 +9380,7 @@ def kernel_chain_dv2(
state_input_ready_l_mbar,
0,
)
update_ready_arrive(update_ready_mbar)
update_ready_arrive(update_ready_mbar.subview(chunk % 2))


@cute.jit
Expand Down
Loading