From 977af0e511a0c9df927f4230cae3b3d43945f9d7 Mon Sep 17 00:00:00 2001 From: Damien Lejeune Date: Thu, 23 Apr 2026 15:22:09 +0000 Subject: [PATCH] Add max len k to UA argument structure --- .../ck_tile/42_unified_attention/example_unified_attention.cpp | 1 + example/ck_tile/42_unified_attention/unified_attention.cpp | 3 ++- example/ck_tile/42_unified_attention/unified_attention.hpp | 3 ++- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/example/ck_tile/42_unified_attention/example_unified_attention.cpp b/example/ck_tile/42_unified_attention/example_unified_attention.cpp index 03e5697ba0..b1d8e903f6 100644 --- a/example/ck_tile/42_unified_attention/example_unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/example_unified_attention.cpp @@ -420,6 +420,7 @@ bool run_impl(const Problem& problem, const RunConfig& run_config) }; ck_tile::index_t max_kv_len = max_element(eff_kv_lens); + args.max_seqlen_k = max_kv_len; ck_tile::index_t max_num_blocks_per_seq = (max_kv_len + problem.page_blk_size - 1) / problem.page_blk_size; diff --git a/example/ck_tile/42_unified_attention/unified_attention.cpp b/example/ck_tile/42_unified_attention/unified_attention.cpp index 7bb33f8dd0..35e8326127 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -80,7 +80,8 @@ std::ostream& operator<<(std::ostream& stream, const unified_attention_args& arg // stream << ", query_start_len_ptr="; // write_ptr(stream, static_cast(args.query_start_len_ptr)); return stream << ", num_seqs=" << args.num_seqs - << ", max_seqlen_q=" << args.max_seqlen_q << " }"; + << ", max_seqlen_q=" << args.max_seqlen_q + << ", max_seqlen_k=" << args.max_seqlen_k << " }"; } // Helper macro to reduce dispatch boilerplate. diff --git a/example/ck_tile/42_unified_attention/unified_attention.hpp b/example/ck_tile/42_unified_attention/unified_attention.hpp index 7bf52d09a6..bdff0a2c33 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.hpp +++ b/example/ck_tile/42_unified_attention/unified_attention.hpp @@ -66,7 +66,8 @@ struct unified_attention_args const int32_t* query_start_len_ptr; // [num_seqs+1] index_t num_seqs; // number of batches for q - index_t max_seqlen_q = 0; // max query length across all batches (0 = unknown) + index_t max_seqlen_q = 0; // max query length across all batches (0 = unknown) + index_t max_seqlen_k = 0; // max KV length across seqs in seq_lens (0 = unknown / not set) }; std::ostream& operator<<(std::ostream& stream,