diff --git a/3rdparty/vendor_patches/flashinfer-kda-bt16.patch b/3rdparty/vendor_patches/flashinfer-kda-bt16.patch index 5941dc6f16e7..e5463bc28a15 100644 --- a/3rdparty/vendor_patches/flashinfer-kda-bt16.patch +++ b/3rdparty/vendor_patches/flashinfer-kda-bt16.patch @@ -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 diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 59503794c989..87e7e9aff20f 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -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 diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py index 38903d429cbc..f90bf435e2d1 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_chunked_bt16.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/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