Skip to content

Repining CK to f33252ce and fix build/test errors #4384

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
JiaLuo-CAN wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Repining CK to f33252ce and fix build/test errors #4384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1879c46
repin CK to f33252ce (rocm-libraries #9740)
JiaLuo-CAN Jul 25, 2026
7bd78c3
fix BlockSize type to align with CK
JiaLuo-CAN Jun 16, 2026
75ab129
Try fix with kHasMask and kHasSink (mha_varlen_fwd_kernels.cu)
JiaLuo-CAN Jul 12, 2026
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
2 changes: 1 addition & 1 deletion 3rdparty/composable_kernel
Open in desktop
Submodule composable_kernel updated 3039 files
2 changes: 1 addition & 1 deletion csrc/ck_deepgemm/include/deepgemm_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ void grouped_flatmm(KernelArguments& args, ck_stream_config& s)
auto kargs = Kernel::MakeKernelArgs(args);

const dim3 grids = Kernel::GridSize(kargs);
constexpr dim3 blocks = Kernel::BlockSize();
const dim3 blocks = Kernel::BlockSize();

ck_tile::launch_kernel(
s, ck_tile::make_kernel<FlatmmConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
Expand Down
2 changes: 1 addition & 1 deletion csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -246,7 +246,7 @@ void moe_gemm(const MoeFlatmmHostArgs& args, const ck_stream_config& s)
auto kargs = Kernel::MakeKernelArgs(args);

const dim3 grids = Kernel::GridSize(kargs);
constexpr dim3 blocks = Kernel::BlockSize();
const dim3 blocks = Kernel::BlockSize();

// if(!Kernel::IsSupportedArgument(kargs))
// {
Expand Down
2 changes: 1 addition & 1 deletion csrc/cktile_gemm_a8w8_bpreshuffle/include/gemm_a8w8_bpreshuffle_cktile_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ float flatmm_calc(const ck_tile::ScaleFlatmmHostArgs<ScaleM, ScaleN>& args,
auto kargs = Kernel::MakeKernelArgs(args);

const dim3 grids = Kernel::GridSize(kargs);
constexpr dim3 blocks = Kernel::BlockSize();
const dim3 blocks = Kernel::BlockSize();

if(!Kernel::IsSupportedArgument(kargs))
{
Expand Down
2 changes: 1 addition & 1 deletion csrc/cpp_itfs/mha_fwd_batch_prefill.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ float mha_batch_prefill(mha_batch_prefill_args args,
int head_size_q = args.hdim_q;
int head_size_v = args.hdim_v;
bool has_dropout = args.p_drop > 0.f;
bool has_sink = args.sink_size > 0;
bool has_sink = args.sink_size > 0 || args.sink_ptr != nullptr;

// The kUseGlobalLoad decision (>2GB KV cache → use `global_load_lds_*`
// instead of SRD `buffer_load_*`) is made per-arm inside the auto-generated
Expand Down
2 changes: 1 addition & 1 deletion csrc/py_itfs_ck/mha_fwd_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ mha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
static_cast<int>(bias_type),
has_lse,
static_cast<int>(qscale_type),
mask.sink > 0, // has_sink
(mask.sink > 0) || (sink_ptr != nullptr), // has_sink: true for streaming-sink window OR a learned per-head sink_ptr (e.g. gpt-oss, sink_size==0)
q.data_ptr(),
k.data_ptr(),
v.data_ptr(),
Expand Down
4 changes: 2 additions & 2 deletions csrc/py_itfs_ck/mha_varlen_fwd_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ mha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
static_cast<int>(bias_type),
has_lse,
static_cast<int>(qscale_type),
mask.sink > 0, // hsa_sink
(mask.sink > 0) || (sink_ptr != nullptr), // has_sink: true for streaming-sink window OR a learned per-head sink_ptr (e.g. gpt-oss, sink_size==0)
q.data_ptr(),
k.data_ptr(),
v.data_ptr(),
Expand Down Expand Up @@ -496,7 +496,7 @@ mha_varlen_fwd(
std::string mask_identify = "b:" + std::to_string(window_size_left) + "," + std::to_string(window_size_right) + "," + std::to_string(sink_size);
mask = mask_info::decode(mask_identify, max_seqlen_q, max_seqlen_k); // local
}
bool has_sink = mask.sink > 0;
bool has_sink = (mask.sink > 0) || sink_ptr.has_value();
CHECK_SHAPE(q, total_q, num_heads, head_size_q);
if (!paged_KV) {
const int total_k = k.size(0);
Expand Down
Loading