Add jenga test and pre-commit

example/ck_tile/50_sparse_attn/CMakeLists.txt

@@ -6,7 +6,7 @@
 set(INST_TARGETS ${SUPPORTED_GPU_TARGETS})
 set(GPU_TARGETS ${SUPPORTED_GPU_TARGETS})
 
-message(STATUS "VSA Sparse Attention: SUPPORTED_GPU_TARGETS=${SUPPORTED_GPU_TARGETS}, INST_TARGETS=${INST_TARGETS}")
+message(STATUS "Sparse Attention: SUPPORTED_GPU_TARGETS=${SUPPORTED_GPU_TARGETS}, INST_TARGETS=${INST_TARGETS}")
 
 list(FILTER INST_TARGETS INCLUDE REGEX "gfx9|gfx12")
 if(NOT INST_TARGETS)
@@ -14,7 +14,7 @@ if(NOT INST_TARGETS)
   return()
 endif()
 
-message(STATUS "Building VSA Sparse Attention for targets: ${INST_TARGETS}")
+message(STATUS "Building Sparse Attention (Jenga & VSA) for targets: ${INST_TARGETS}")
 
 # Code generation scripts
 file(GLOB_RECURSE CODE_GEN_SCRIPTS CONFIGURE_DEPENDS
@@ -23,14 +23,81 @@ file(GLOB_RECURSE CODE_GEN_SCRIPTS CONFIGURE_DEPENDS
 )
 set_directory_properties(PROPERTIES CMAKE_CONFIGURE_DEPENDS "${CODE_GEN_SCRIPTS}")
 
-# Code generation for VSA (receipt 600 for aiter integration)
+# ============================================================================
+# Jenga Sparse Attention
+# ============================================================================
+set(SPARSE_ATTN_JENGA_CODE_GEN_ARGS
+  ${CMAKE_CURRENT_LIST_DIR}/generate.py
+  --api fwd_jenga
+  --receipt 600
+)
+
+# Generate list of Jenga kernels (at configure time, only list)
+execute_process(
+  COMMAND ${Python3_EXECUTABLE} ${SPARSE_ATTN_JENGA_CODE_GEN_ARGS}
+  --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/jenga_blob_list.txt
+  RESULT_VARIABLE ret
+)
+if(ret AND NOT ret EQUAL 0)
+  message(FATAL_ERROR "Failed to generate Jenga kernel list")
+endif()
+
+file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/jenga_blob_list.txt SPARSE_ATTN_JENGA_GEN_BLOBS)
+
+# Generate Jenga kernel source files at build time
+add_custom_command(
+  OUTPUT ${SPARSE_ATTN_JENGA_GEN_BLOBS}
+  COMMAND ${Python3_EXECUTABLE} ${SPARSE_ATTN_JENGA_CODE_GEN_ARGS}
+  --output_dir ${CMAKE_CURRENT_BINARY_DIR}
+  DEPENDS ${CODE_GEN_SCRIPTS}
+  COMMENT "Generate CK Tile Jenga Sparse Attention kernels"
+)
+
+message(STATUS "Jenga kernel files to be generated: ${SPARSE_ATTN_JENGA_GEN_BLOBS}")
+
+# Jenga Instances
+set(SPARSE_ATTN_JENGA_INSTANCES "tile_sparse_attn_jenga_instances")
+
+add_library(${SPARSE_ATTN_JENGA_INSTANCES} OBJECT EXCLUDE_FROM_ALL
+  ${SPARSE_ATTN_JENGA_GEN_BLOBS}
+  ${CMAKE_CURRENT_LIST_DIR}/jenga_sparse_attention.cu
+)
+target_include_directories(${SPARSE_ATTN_JENGA_INSTANCES} PRIVATE 
+  ${CMAKE_CURRENT_LIST_DIR}
+  ${PROJECT_SOURCE_DIR}/include/ck_tile/ops/sparse_attn
+)
+set_source_files_properties(${SPARSE_ATTN_JENGA_GEN_BLOBS} PROPERTIES LANGUAGE HIP)
+set_source_files_properties(${CMAKE_CURRENT_LIST_DIR}/jenga_sparse_attention.cu PROPERTIES LANGUAGE HIP)
+set_property(TARGET ${SPARSE_ATTN_JENGA_INSTANCES} PROPERTY HIP_ARCHITECTURES ${INST_TARGETS})
+
+target_compile_options(${SPARSE_ATTN_JENGA_INSTANCES} PRIVATE
+  -DCK_TILE_USE_BUFFER_ADDRESSING_BUILTIN
+  -DCK_TILE_FMHA_FWD_FAST_EXP2
+  -Wno-undefined-func-template
+  -Wno-float-equal
+)
+
+# Jenga Example executable
+set(EXAMPLE_JENGA_SPARSE_ATTN "tile_example_jenga_sparse_attn")
+message(DEBUG "adding example ${EXAMPLE_JENGA_SPARSE_ATTN}")
+add_executable(${EXAMPLE_JENGA_SPARSE_ATTN} EXCLUDE_FROM_ALL test_jenga_sparse_attn.cpp)
+target_link_libraries(${EXAMPLE_JENGA_SPARSE_ATTN} ${SPARSE_ATTN_JENGA_INSTANCES})
+target_include_directories(${EXAMPLE_JENGA_SPARSE_ATTN} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+target_compile_options(${EXAMPLE_JENGA_SPARSE_ATTN} PRIVATE
+  -Wno-undefined-func-template
+  -Wno-float-equal
+)
+
+# ============================================================================
+# VSA Sparse Attention
+# ============================================================================
 set(SPARSE_ATTN_VSA_CODE_GEN_ARGS
   ${CMAKE_CURRENT_LIST_DIR}/generate.py
   --api fwd_vsa
   --receipt 600
 )
 
-# Generate list of VSA kernels (at configure time, only list, not generate)
+# Generate list of VSA kernels (at configure time, only list)
 execute_process(
   COMMAND ${Python3_EXECUTABLE} ${SPARSE_ATTN_VSA_CODE_GEN_ARGS}
   --list_blobs ${CMAKE_CURRENT_BINARY_DIR}/vsa_blob_list.txt
@@ -42,7 +109,7 @@ endif()
 
 file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/vsa_blob_list.txt SPARSE_ATTN_VSA_GEN_BLOBS)
 
-# Generate the kernel source files at build time (not configure time)
+# Generate VSA kernel source files at build time
 add_custom_command(
   OUTPUT ${SPARSE_ATTN_VSA_GEN_BLOBS}
   COMMAND ${Python3_EXECUTABLE} ${SPARSE_ATTN_VSA_CODE_GEN_ARGS}
@@ -68,7 +135,6 @@ set_source_files_properties(${SPARSE_ATTN_VSA_GEN_BLOBS} PROPERTIES LANGUAGE HIP
 set_source_files_properties(${CMAKE_CURRENT_LIST_DIR}/vsa_sparse_attention.cu PROPERTIES LANGUAGE HIP)
 set_property(TARGET ${SPARSE_ATTN_VSA_INSTANCES} PROPERTY HIP_ARCHITECTURES ${INST_TARGETS})
 
-# Compile options
 target_compile_options(${SPARSE_ATTN_VSA_INSTANCES} PRIVATE
   -DCK_TILE_USE_BUFFER_ADDRESSING_BUILTIN
   -DCK_TILE_FMHA_FWD_FAST_EXP2
@@ -76,12 +142,13 @@ target_compile_options(${SPARSE_ATTN_VSA_INSTANCES} PRIVATE
   -Wno-float-equal
 )
 
-# Test executable
-set(TEST_VSA_SPARSE_ATTN "tile_test_vsa_sparse_attn")
-add_executable(${TEST_VSA_SPARSE_ATTN} EXCLUDE_FROM_ALL test_vsa_sparse_attn.cpp)
-target_link_libraries(${TEST_VSA_SPARSE_ATTN} ${SPARSE_ATTN_VSA_INSTANCES})
-target_include_directories(${TEST_VSA_SPARSE_ATTN} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
-target_compile_options(${TEST_VSA_SPARSE_ATTN} PRIVATE
+# VSA Example executable
+set(EXAMPLE_VSA_SPARSE_ATTN "tile_example_vsa_sparse_attn")
+message(DEBUG "adding example ${EXAMPLE_VSA_SPARSE_ATTN}")
+add_executable(${EXAMPLE_VSA_SPARSE_ATTN} EXCLUDE_FROM_ALL test_vsa_sparse_attn.cpp)
+target_link_libraries(${EXAMPLE_VSA_SPARSE_ATTN} ${SPARSE_ATTN_VSA_INSTANCES})
+target_include_directories(${EXAMPLE_VSA_SPARSE_ATTN} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+target_compile_options(${EXAMPLE_VSA_SPARSE_ATTN} PRIVATE
   -Wno-undefined-func-template
   -Wno-float-equal
 )

example/ck_tile/50_sparse_attn/codegen/ops/fmha_fwd_jenga.py

@@ -20,6 +20,7 @@ from codegen.cpp_symbol_map import (
     MODE_MAP,
     PIPELINE_ENUM_MAP,
     PIPELINE_MAP,
+    SQUANT_MAP,
     get_mask_check_map,
     get_mask_map,
 )
@@ -78,7 +79,7 @@ using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad},
                                                     false,
                                                     {F_lse},
                                                     {F_dropout},
-                                                    {F_squant},
+                                                    {F_squant_enum},
                                                     {F_occupancy},
                                                     {F_skip}>;
 
@@ -596,6 +597,7 @@ class FmhaFwdKernel:
             F_lse=BOOL_MAP[self.F_pipeline.F_lse],
             F_dropout=BOOL_MAP[self.F_pipeline.F_dropout],
             F_squant=BOOL_MAP[self.F_pipeline.F_squant],
+            F_squant_enum=SQUANT_MAP[self.F_pipeline.F_squant],
             F_skip=BOOL_MAP[self.F_pipeline.F_skip],
             F_occupancy=self.F_tile.F_occupancy,
             F_pipeline_enum=PIPELINE_ENUM_MAP[self.F_pipeline.tag],

example/ck_tile/50_sparse_attn/jenga_sparse_attention.cu

@@ -5,6 +5,7 @@
 #include "fmha_fwd_trek.hpp"
 #include "ck_tile/core.hpp"
 #include "ck_tile/host/host_tensor.hpp"
+#include "ck_tile/host/device_memory.hpp"
 
 ck_tile::HostTensor<DataType>
 jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
@@ -12,36 +13,32 @@ jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
                        ck_tile::HostTensor<DataType>& TV,
                        ck_tile::HostTensor<DataType>& Tblock_relation_onehot,
                        ck_tile::HostTensor<DataType>& Y,
-                       std::optional<ck_tile::HostTensor<DataType>> bias       = std::nullopt,
-                       std::optional<ck_tile::HostTensor<DataType>> lse        = std::nullopt,
-                       std::optional<ck_tile::HostTensor<DataType>> seqstart_q = std::nullopt,
-                       std::optional<ck_tile::HostTensor<DataType>> seqstart_k = std::nullopt,
-                       int bias_type                                           = 0,
-                       int batch                                               = 0,
-                       int nhead                                               = 0,
-                       int nhead_k                                             = 0,
-                       int seqlen_q                                            = 0,
-                       int seqlen_k                                            = 0,
-                       int hdim_q                                              = 0,
-                       int hdim_v                                              = 0,
-                       int mode                                                = 0,
-                       bool i_perm                                             = true,
-                       bool o_perm                                             = true,
-                       int max_seqlen_q                                        = 0,
-                       int max_seqlen_k                                        = 0)
+                       std::optional<ck_tile::HostTensor<DataType>> bias,
+                       std::optional<ck_tile::HostTensor<DataType>> lse,
+                       std::optional<ck_tile::HostTensor<DataType>> seqstart_q,
+                       std::optional<ck_tile::HostTensor<DataType>> seqstart_k,
+                       int bias_type,
+                       int batch,
+                       int nhead,
+                       int nhead_k,
+                       int seqlen_q,
+                       int seqlen_k,
+                       int hdim_q,
+                       int hdim_v,
+                       int mode,
+                       bool i_perm,
+                       bool o_perm,
+                       int max_seqlen_q,
+                       int max_seqlen_k)
 {
     std::string data_type = "fp16";
-    if(TQ.dtype() == ck_tile::bf16_t)
-    {
-        data_type = "bf16";
-    }
+    // DataType is determined at compile time via template
 
     if(max_seqlen_q == 0)
         max_seqlen_q = seqlen_q;
     if(max_seqlen_k == 0)
         max_seqlen_k = seqlen_k;
     bool is_v_rowmajor          = true;
-    int seqlen_knew             = 0;
     float scale_s               = 1.0 / ck_tile::sqrt(static_cast<float>(hdim_q));
     float scale_p               = 1.f;
     float scale_o               = 1.f;
@@ -50,7 +47,6 @@ jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
     std::string msk_str = "0";
     mask_info mask      = mask_info::decode(msk_str, seqlen_q, seqlen_k);
 
-    const ck_tile::index_t shape_batch    = (mode == 0 ? batch : 1);
     const ck_tile::index_t shape_seqlen_q = (mode == 0 ? seqlen_q : max_seqlen_q);
     const ck_tile::index_t shape_seqlen_k = (mode == 0 ? seqlen_k : max_seqlen_k);
 
@@ -61,71 +57,74 @@ jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
                                          1,
                                          false};
 
+    // Create device memory and copy data to device
+    ck_tile::DeviceMem q_buf(TQ.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem k_buf(TK.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem v_buf(TV.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem block_relation_buf(Tblock_relation_onehot.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem o_buf(Y.get_element_space_size_in_bytes());
+
+    q_buf.ToDevice(TQ.data());
+    k_buf.ToDevice(TK.data());
+    v_buf.ToDevice(TV.data());
+    block_relation_buf.ToDevice(Tblock_relation_onehot.data());
+
+    // Optional buffers
+    ck_tile::DeviceMem bias_buf(bias ? bias->get_element_space_size_in_bytes() : 0);
+    ck_tile::DeviceMem lse_buf(lse ? lse->get_element_space_size_in_bytes() : 0);
+    ck_tile::DeviceMem seqstart_q_buf(seqstart_q ? seqstart_q->get_element_space_size_in_bytes() : 0);
+    ck_tile::DeviceMem seqstart_k_buf(seqstart_k ? seqstart_k->get_element_space_size_in_bytes() : 0);
+
+    if(bias)
+        bias_buf.ToDevice(bias->data());
+    if(lse)
+        lse_buf.ToDevice(lse->data());
+    if(seqstart_q)
+        seqstart_q_buf.ToDevice(seqstart_q->data());
+    if(seqstart_k)
+        seqstart_k_buf.ToDevice(seqstart_k->data());
+
     const auto init_args = [&](auto& args) {
         assert(nhead % nhead_k == 0);
-        const ck_tile::index_t stride_q    = (i_perm ? hdim_q : nhead * hdim_q);
-        const ck_tile::index_t stride_k    = (i_perm ? hdim_q : nhead_k * hdim_q);
-        const ck_tile::index_t stride_knew = (i_perm ? hdim_q : nhead_k * hdim_q);
-        const ck_tile::index_t stride_v    = [&]() {
+        const ck_tile::index_t stride_q = (i_perm ? hdim_q : nhead * hdim_q);
+        const ck_tile::index_t stride_k = (i_perm ? hdim_q : nhead_k * hdim_q);
+        const ck_tile::index_t stride_v = [&]() {
             if(is_v_rowmajor)
                 return i_perm ? hdim_v : nhead_k * hdim_v;
             else
                 return (i_perm ? shape_seqlen_k : nhead_k * shape_seqlen_k);
         }();
-        const ck_tile::index_t stride_vnew = [&]() {
-            if(is_v_rowmajor)
-                return i_perm ? hdim_v : nhead_k * hdim_v;
-            else
-                return i_perm ? seqlen_knew : nhead_k * seqlen_knew;
-        }();
         const ck_tile::index_t stride_bias    = (i_perm ? max_seqlen_k : 1 * max_seqlen_k);
         const ck_tile::index_t stride_randval = (max_seqlen_k);
-        const ck_tile::index_t stride_o_acc   = (hdim_v);
         const ck_tile::index_t stride_o       = (o_perm ? hdim_v : nhead * hdim_v);
         // setup nhead_stride_* arguments
-        const ck_tile::index_t nhead_stride_q    = (i_perm ? shape_seqlen_q * hdim_q : hdim_q);
-        const ck_tile::index_t nhead_stride_k    = i_perm ? shape_seqlen_k * hdim_q : hdim_q;
-        const ck_tile::index_t nhead_stride_knew = (i_perm ? seqlen_knew * hdim_q : hdim_q);
-        const ck_tile::index_t nhead_stride_v    = [&]() {
+        const ck_tile::index_t nhead_stride_q = (i_perm ? shape_seqlen_q * hdim_q : hdim_q);
+        const ck_tile::index_t nhead_stride_k = i_perm ? shape_seqlen_k * hdim_q : hdim_q;
+        const ck_tile::index_t nhead_stride_v = [&]() {
             if(is_v_rowmajor)
                 return i_perm ? shape_seqlen_k * hdim_v : hdim_v;
             else
                 return i_perm ? hdim_v * shape_seqlen_k : shape_seqlen_k;
         }();
-        const ck_tile::index_t nhead_stride_vnew = [&]() {
-            if(is_v_rowmajor)
-                return i_perm ? seqlen_knew * hdim_v : hdim_v;
-            else
-                return i_perm ? hdim_v * seqlen_knew : seqlen_knew;
-        }();
         const ck_tile::index_t nhead_stride_bias =
             (i_perm ? 0 * shape_seqlen_q * max_seqlen_k : 0 * max_seqlen_k);
         const ck_tile::index_t nhead_stride_randval = (shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t nhead_stride_lse     = shape_seqlen_q;
-        const ck_tile::index_t nhead_stride_lse_acc = (shape_seqlen_q);
-        const ck_tile::index_t nhead_stride_o_acc   = (shape_seqlen_q * hdim_v);
         const ck_tile::index_t nhead_stride_o       = (o_perm ? shape_seqlen_q * hdim_v : hdim_v);
         // setup batch_stride_* arguments
         const ck_tile::index_t batch_stride_q       = (nhead * shape_seqlen_q * hdim_q);
         const ck_tile::index_t batch_stride_k       = nhead_k * shape_seqlen_k * hdim_q;
-        const ck_tile::index_t batch_stride_knew    = (nhead_k * seqlen_knew * hdim_q);
         const ck_tile::index_t batch_stride_v       = nhead_k * hdim_v * shape_seqlen_k;
-        const ck_tile::index_t batch_stride_vnew    = (nhead_k * hdim_v * seqlen_knew);
         const ck_tile::index_t batch_stride_bias    = (0 * nhead * shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t batch_stride_randval = (nhead * shape_seqlen_q * max_seqlen_k);
         const ck_tile::index_t batch_stride_lse     = (nhead * shape_seqlen_q);
-        const ck_tile::index_t batch_stride_lse_acc = (nhead * shape_seqlen_q);
-        const ck_tile::index_t batch_stride_o_acc   = (nhead * shape_seqlen_q * hdim_v);
         const ck_tile::index_t batch_stride_o       = (nhead * shape_seqlen_q * hdim_v);
-        // const ck_tile::index_t batch_stride_block_table = (max_num_page_blocks / batch);
-        // setup split_stride_* arguments (only used in split-kv kernel)
-        const ck_tile::index_t split_stride_lse_acc = (shape_seqlen_q);
-        const ck_tile::index_t split_stride_o_acc   = (shape_seqlen_q * hdim_v);
 
-        args.q_ptr                     = TQ.data_ptr();
-        args.k_ptr                     = TK.data_ptr();
-        args.v_ptr                     = TV.data_ptr();
-        args.block_relation_onehot_ptr = Tblock_relation_onehot.data_ptr();
+        // Use device buffer pointers instead of host tensor data pointers
+        args.q_ptr                     = q_buf.GetDeviceBuffer();
+        args.k_ptr                     = k_buf.GetDeviceBuffer();
+        args.v_ptr                     = v_buf.GetDeviceBuffer();
+        args.block_relation_onehot_ptr = block_relation_buf.GetDeviceBuffer();
 
         args.batch    = batch;
         args.seqlen_q = shape_seqlen_q; // unused in group mode
@@ -144,14 +143,12 @@ jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
         args.batch_stride_k = batch_stride_k;
         args.batch_stride_v = batch_stride_v;
 
-        // args.bias_ptr = bias.type == bias_enum::alibi ? alibi_slope_buf.GetDeviceBuffer()
-        //                                                 : bias_buf.GetDeviceBuffer();
-        args.bias_ptr = bias ? bias->data_ptr() : nullptr;
-        args.lse_ptr  = lse ? lse->data_ptr() : nullptr;
-        args.o_ptr    = Y.data_ptr();
+        args.bias_ptr = bias ? bias_buf.GetDeviceBuffer() : nullptr;
+        args.lse_ptr  = lse ? lse_buf.GetDeviceBuffer() : nullptr;
+        args.o_ptr    = o_buf.GetDeviceBuffer();
 
-        args.seqstart_q_ptr = (mode == 1 ? seqstart_q->data_ptr() : nullptr);
-        args.seqstart_k_ptr = (mode == 1 ? seqstart_k->data_ptr() : nullptr);
+        args.seqstart_q_ptr = (mode == 1 ? seqstart_q_buf.GetDeviceBuffer() : nullptr);
+        args.seqstart_k_ptr = (mode == 1 ? seqstart_k_buf.GetDeviceBuffer() : nullptr);
         args.seqlen_k_ptr   = nullptr;
 
         args.seqlen_k     = shape_seqlen_k; // unused in group mode (or kvcache enabled)
@@ -210,5 +207,8 @@ jenga_sparse_attention(ck_tile::HostTensor<DataType>& TQ,
 
     fmha_jenga_fwd(fmha_traits, args, stream_config);
 
+    // Copy output back to host
+    Y = o_buf.ToHost<DataType>();
+
     return Y;
 }

example/ck_tile/50_sparse_attn/test_jenga_sparse_attn.cpp

@@ -0,0 +1,455 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+//
+// Test for jenga_sparse_attention function
+
+#include <iostream>
+#include <vector>
+#include <cmath>
+#include <random>
+#include <string>
+#include <algorithm>
+#include <numeric>
+#include <chrono>
+
+#include "ck_tile/host.hpp"
+#include "ck_tile/core.hpp"
+
+#include "jenga_sparse_attention.h"
+
+// ============================================================================
+// Helper Functions
+// ============================================================================
+
+// Reference implementation: blocked attention
+template <typename T, typename AccT = float>
+void reference_blocked_attention(
+    const ck_tile::HostTensor<T>& q,              // [B, H, S_q, D]
+    const ck_tile::HostTensor<T>& k,              // [B, H, S_k, D]
+    const ck_tile::HostTensor<T>& v,              // [B, H, S_k, D_v]
+    const ck_tile::HostTensor<T>& block_relation, // [B, H, Q_blocks, K_blocks]
+    const ck_tile::HostTensor<T>& bias,           // [B, H, S_q, S_k]
+    ck_tile::HostTensor<T>& output,               // [B, H, S_q, D_v]
+    ck_tile::index_t BLKQ,
+    ck_tile::index_t BLKK,
+    AccT scale)
+{
+    auto q_lengths            = q.get_lengths();
+    ck_tile::index_t batch    = q_lengths[0];
+    ck_tile::index_t nhead    = q_lengths[1];
+    ck_tile::index_t seqlen_q = q_lengths[2];
+    ck_tile::index_t hdim     = q_lengths[3];
+
+    auto v_lengths            = v.get_lengths();
+    ck_tile::index_t seqlen_k = v_lengths[2];
+    ck_tile::index_t hdim_v   = v_lengths[3];
+
+    ck_tile::index_t num_q_blocks = seqlen_q / BLKQ;
+    ck_tile::index_t num_k_blocks = seqlen_k / BLKK;
+
+    for(ck_tile::index_t b = 0; b < batch; ++b)
+    {
+        for(ck_tile::index_t h = 0; h < nhead; ++h)
+        {
+            for(ck_tile::index_t qb = 0; qb < num_q_blocks; ++qb)
+            {
+                ck_tile::index_t q_start = qb * BLKQ;
+                ck_tile::index_t q_end   = q_start + BLKQ;
+
+                // Find relevant K blocks
+                std::vector<ck_tile::index_t> relevant_k_indices;
+                for(ck_tile::index_t kb = 0; kb < num_k_blocks; ++kb)
+                {
+                    if(static_cast<float>(block_relation(b, h, qb, kb)) > 0.5f)
+                    {
+                        relevant_k_indices.push_back(kb);
+                    }
+                }
+
+                if(relevant_k_indices.empty())
+                    continue;
+
+                // For each query position in the block
+                for(ck_tile::index_t sq = q_start; sq < q_end; ++sq)
+                {
+                    std::vector<AccT> scores;
+                    AccT max_score = -std::numeric_limits<AccT>::infinity();
+
+                    for(auto kb : relevant_k_indices)
+                    {
+                        ck_tile::index_t k_start = kb * BLKK;
+                        ck_tile::index_t k_end   = k_start + BLKK;
+
+                        for(ck_tile::index_t sk = k_start; sk < k_end; ++sk)
+                        {
+                            AccT score = 0.0f;
+                            for(ck_tile::index_t d = 0; d < hdim; ++d)
+                            {
+                                score += static_cast<AccT>(q(b, h, sq, d)) *
+                                         static_cast<AccT>(k(b, h, sk, d));
+                            }
+                            score = score * scale + static_cast<AccT>(bias(b, h, sq, sk));
+                            scores.push_back(score);
+                            max_score = std::max(max_score, score);
+                        }
+                    }
+
+                    // Softmax
+                    AccT sum_exp = 0.0f;
+                    for(auto& s : scores)
+                    {
+                        s = std::exp(s - max_score);
+                        sum_exp += s;
+                    }
+                    for(auto& s : scores)
+                    {
+                        s /= sum_exp;
+                    }
+
+                    // Compute output: P @ V
+                    for(ck_tile::index_t dv = 0; dv < hdim_v; ++dv)
+                    {
+                        AccT out_val     = 0.0f;
+                        size_t score_idx = 0;
+
+                        for(auto kb : relevant_k_indices)
+                        {
+                            ck_tile::index_t k_start = kb * BLKK;
+                            ck_tile::index_t k_end   = k_start + BLKK;
+
+                            for(ck_tile::index_t sk = k_start; sk < k_end; ++sk)
+                            {
+                                out_val += scores[score_idx] * static_cast<AccT>(v(b, h, sk, dv));
+                                score_idx++;
+                            }
+                        }
+                        output(b, h, sq, dv) = static_cast<T>(out_val);
+                    }
+                }
+            }
+        }
+    }
+}
+
+// ============================================================================
+// Command line argument parser
+// ============================================================================
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("v", "1", "0:no validation, 1:cpu validation")
+        .insert("mode", "0", "kernel mode. 0:batch, 1:group")
+        .insert("b", "1", "batch size")
+        .insert("h", "4", "num of head for q")
+        .insert("h_k", "-1", "num of head for k/v, -1 means equal to h")
+        .insert("s", "4096", "seqlen_q")
+        .insert("s_k", "-1", "seqlen_k, -1 means equal to s")
+        .insert("d", "128", "head dim for q, k")
+        .insert("d_v", "-1", "head dim for v, -1 means equal to d")
+        .insert("block_size", "128", "block size for sparse attention (BLKQ=BLKK)")
+        .insert("sparsity", "0.5", "sparsity ratio (0.0 = dense, 1.0 = fully sparse)")
+        .insert("iperm", "1", "permute input, 1: b*h*s*d, 0: b*s*h*d")
+        .insert("operm", "1", "permute output")
+        .insert("bias", "0", "bias type: 0:no bias, 1:elementwise, 2:alibi")
+        .insert("lse", "0", "0:not store lse, 1:store lse")
+        .insert("seed", "42", "random seed")
+        .insert("warmup", "5", "warmup iterations")
+        .insert("repeat", "20", "benchmark iterations");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+// ============================================================================
+// Main Test Function
+// ============================================================================
+bool run_test(const ck_tile::ArgParser& arg_parser)
+{
+    using T = DataType; // Use DataType defined in header (half_t)
+
+    // Parse arguments
+    int do_validation           = arg_parser.get_int("v");
+    int mode                    = arg_parser.get_int("mode");
+    ck_tile::index_t batch      = arg_parser.get_int("b");
+    ck_tile::index_t nhead      = arg_parser.get_int("h");
+    ck_tile::index_t nhead_k    = arg_parser.get_int("h_k");
+    ck_tile::index_t seqlen_q   = arg_parser.get_int("s");
+    ck_tile::index_t seqlen_k   = arg_parser.get_int("s_k");
+    ck_tile::index_t hdim_q     = arg_parser.get_int("d");
+    ck_tile::index_t hdim_v     = arg_parser.get_int("d_v");
+    ck_tile::index_t block_size = arg_parser.get_int("block_size");
+    float sparsity              = arg_parser.get_float("sparsity");
+    bool i_perm                 = arg_parser.get_bool("iperm");
+    bool o_perm                 = arg_parser.get_bool("operm");
+    int bias_type               = arg_parser.get_int("bias");
+    bool store_lse              = arg_parser.get_bool("lse");
+    uint32_t seed               = arg_parser.get_uint32("seed");
+    int warmup                  = arg_parser.get_int("warmup");
+    int repeat                  = arg_parser.get_int("repeat");
+
+    // Handle default values
+    if(nhead_k < 0)
+        nhead_k = nhead;
+    if(seqlen_k < 0)
+        seqlen_k = seqlen_q;
+    if(hdim_v < 0)
+        hdim_v = hdim_q;
+
+    ck_tile::index_t BLKQ = block_size;
+    ck_tile::index_t BLKK = block_size;
+
+    // Calculate number of Q and K blocks
+    ck_tile::index_t num_q_blocks = seqlen_q / BLKQ;
+    ck_tile::index_t num_k_blocks = seqlen_k / BLKK;
+
+    std::cout << "============================================================" << std::endl;
+    std::cout << "[Jenga Sparse Attention Test]" << std::endl;
+    std::cout << "============================================================" << std::endl;
+    std::cout << "  Batch: " << batch << ", nhead_q: " << nhead << ", nhead_k: " << nhead_k
+              << std::endl;
+    std::cout << "  seqlen_q: " << seqlen_q << ", seqlen_k: " << seqlen_k << std::endl;
+    std::cout << "  hdim_q: " << hdim_q << ", hdim_v: " << hdim_v << std::endl;
+    std::cout << "  block_size: " << block_size << " (BLKQ=" << BLKQ << ", BLKK=" << BLKK << ")"
+              << std::endl;
+    std::cout << "  num_q_blocks: " << num_q_blocks << ", num_k_blocks: " << num_k_blocks
+              << std::endl;
+    std::cout << "  sparsity: " << sparsity << std::endl;
+    std::cout << "  i_perm: " << i_perm << ", o_perm: " << o_perm << std::endl;
+
+    // Create host tensors (using BHSD layout when i_perm=true)
+    ck_tile::HostTensor<T> q_host({batch, nhead, seqlen_q, hdim_q});
+    ck_tile::HostTensor<T> k_host({batch, nhead_k, seqlen_k, hdim_q});
+    ck_tile::HostTensor<T> v_host({batch, nhead_k, seqlen_k, hdim_v});
+    ck_tile::HostTensor<T> output_host({batch, nhead, seqlen_q, hdim_v});
+    ck_tile::HostTensor<T> output_ref({batch, nhead, seqlen_q, hdim_v});
+
+    // Bias tensor [B, H, S_q, S_k]
+    ck_tile::HostTensor<T> bias_host({batch, nhead, seqlen_q, seqlen_k});
+
+    // Block relation onehot: [B, H, Q_blocks, K_blocks]
+    ck_tile::HostTensor<T> block_relation_onehot({batch, nhead, num_q_blocks, num_k_blocks});
+
+    // LSE tensor (optional)
+    ck_tile::HostTensor<T> lse_host({batch, nhead, seqlen_q});
+
+    // Initialize tensors with random values
+    std::cout << "\nInitializing tensors..." << std::endl;
+    ck_tile::FillUniformDistribution<T>{-0.5f, 0.5f, seed}(q_host);
+    ck_tile::FillUniformDistribution<T>{-0.5f, 0.5f, seed + 1}(k_host);
+    ck_tile::FillUniformDistribution<T>{-0.5f, 0.5f, seed + 2}(v_host);
+
+    // Initialize bias to zero
+    std::fill(bias_host.mData.begin(), bias_host.mData.end(), static_cast<T>(0.0f));
+
+    // Initialize block_relation_onehot with sparse pattern
+    std::mt19937 rng(seed + 100);
+    std::uniform_real_distribution<float> dist(0.0f, 1.0f);
+    ck_tile::index_t total_blocks  = 0;
+    ck_tile::index_t active_blocks = 0;
+
+    for(ck_tile::index_t b = 0; b < batch; ++b)
+    {
+        for(ck_tile::index_t h = 0; h < nhead; ++h)
+        {
+            for(ck_tile::index_t qb = 0; qb < num_q_blocks; ++qb)
+            {
+                for(ck_tile::index_t kb = 0; kb < num_k_blocks; ++kb)
+                {
+                    total_blocks++;
+                    bool is_diagonal   = (qb == kb && qb < num_k_blocks);
+                    bool random_active = (dist(rng) > sparsity);
+
+                    if(is_diagonal || random_active)
+                    {
+                        block_relation_onehot(b, h, qb, kb) = static_cast<T>(1.0f);
+                        active_blocks++;
+                    }
+                    else
+                    {
+                        block_relation_onehot(b, h, qb, kb) = static_cast<T>(0.0f);
+                    }
+                }
+            }
+        }
+    }
+
+    float actual_sparsity =
+        1.0f - static_cast<float>(active_blocks) / static_cast<float>(total_blocks);
+    std::cout << "  Actual sparsity: " << actual_sparsity << " (" << active_blocks << "/"
+              << total_blocks << " blocks active)" << std::endl;
+
+    // Optional tensors
+    std::optional<ck_tile::HostTensor<T>> bias_opt       = std::nullopt;
+    std::optional<ck_tile::HostTensor<T>> lse_opt        = std::nullopt;
+    std::optional<ck_tile::HostTensor<T>> seqstart_q_opt = std::nullopt;
+    std::optional<ck_tile::HostTensor<T>> seqstart_k_opt = std::nullopt;
+
+    if(bias_type != 0)
+    {
+        bias_opt = bias_host;
+    }
+    if(store_lse)
+    {
+        lse_opt = lse_host;
+    }
+
+    // Run kernel
+    std::cout << "\n--- Running Jenga sparse attention kernel ---" << std::endl;
+
+    try
+    {
+        // Warmup
+        for(int i = 0; i < warmup; ++i)
+        {
+            jenga_sparse_attention(q_host,
+                                   k_host,
+                                   v_host,
+                                   block_relation_onehot,
+                                   output_host,
+                                   bias_opt,
+                                   lse_opt,
+                                   seqstart_q_opt,
+                                   seqstart_k_opt,
+                                   bias_type,
+                                   batch,
+                                   nhead,
+                                   nhead_k,
+                                   seqlen_q,
+                                   seqlen_k,
+                                   hdim_q,
+                                   hdim_v,
+                                   mode,
+                                   i_perm,
+                                   o_perm,
+                                   seqlen_q,
+                                   seqlen_k);
+        }
+
+        // Benchmark
+        [[maybe_unused]] auto sync_status1 = hipDeviceSynchronize();
+        auto start                         = std::chrono::high_resolution_clock::now();
+
+        for(int i = 0; i < repeat; ++i)
+        {
+            jenga_sparse_attention(q_host,
+                                   k_host,
+                                   v_host,
+                                   block_relation_onehot,
+                                   output_host,
+                                   bias_opt,
+                                   lse_opt,
+                                   seqstart_q_opt,
+                                   seqstart_k_opt,
+                                   bias_type,
+                                   batch,
+                                   nhead,
+                                   nhead_k,
+                                   seqlen_q,
+                                   seqlen_k,
+                                   hdim_q,
+                                   hdim_v,
+                                   mode,
+                                   i_perm,
+                                   o_perm,
+                                   seqlen_q,
+                                   seqlen_k);
+        }
+
+        [[maybe_unused]] auto sync_status2 = hipDeviceSynchronize();
+        auto end                           = std::chrono::high_resolution_clock::now();
+        double avg_time_ms =
+            std::chrono::duration<double, std::milli>(end - start).count() / repeat;
+
+        std::cout << "\n>>>> Jenga sparse attention average time: " << avg_time_ms << " ms <<<<"
+                  << std::endl;
+    }
+    catch(const std::exception& e)
+    {
+        std::cerr << "Error during kernel execution: " << e.what() << std::endl;
+        return false;
+    }
+
+    // Validation
+    bool pass = true;
+    if(do_validation)
+    {
+        std::cout << "\n--- Performing CPU validation ---" << std::endl;
+
+        float scale = 1.0f / std::sqrt(static_cast<float>(hdim_q));
+
+        std::cout << "Computing reference output..." << std::endl;
+        reference_blocked_attention(q_host,
+                                    k_host,
+                                    v_host,
+                                    block_relation_onehot,
+                                    bias_host,
+                                    output_ref,
+                                    BLKQ,
+                                    BLKK,
+                                    scale);
+
+        // Compare results
+        double rtol = 1e-2;
+        double atol = 4e-2;
+
+        float max_diff     = 0.0f;
+        float max_rel_diff = 0.0f;
+        size_t num_errors  = 0;
+
+        for(size_t i = 0; i < output_host.mData.size(); ++i)
+        {
+            float gpu_val  = static_cast<float>(output_host.mData[i]);
+            float ref_val  = static_cast<float>(output_ref.mData[i]);
+            float diff     = std::abs(gpu_val - ref_val);
+            float rel_diff = (std::abs(ref_val) > 1e-6f) ? diff / std::abs(ref_val) : diff;
+
+            max_diff     = std::max(max_diff, diff);
+            max_rel_diff = std::max(max_rel_diff, rel_diff);
+
+            if(diff > atol && rel_diff > rtol)
+            {
+                num_errors++;
+                if(num_errors <= 5)
+                {
+                    std::cout << "  Mismatch at index " << i << ": GPU=" << gpu_val
+                              << ", Ref=" << ref_val << ", Diff=" << diff << std::endl;
+                }
+            }
+        }
+
+        std::cout << "\nValidation results:" << std::endl;
+        std::cout << "  Max absolute difference: " << max_diff << std::endl;
+        std::cout << "  Max relative difference: " << max_rel_diff << std::endl;
+        std::cout << "  Number of mismatches: " << num_errors << " / " << output_host.mData.size()
+                  << std::endl;
+
+        if(num_errors == 0)
+        {
+            std::cout << "\n>>> VALIDATION PASSED <<<" << std::endl;
+        }
+        else
+        {
+            std::cout << "\n>>> VALIDATION FAILED <<<" << std::endl;
+            pass = false;
+        }
+    }
+
+    std::cout << "\n" << (pass ? "TEST PASSED" : "TEST FAILED") << std::endl;
+    return pass;
+}
+
+// ============================================================================
+// Main
+// ============================================================================
+int main(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+    {
+        std::cerr << "Failed to parse arguments" << std::endl;
+        return -1;
+    }
+
+    bool test_result = run_test(arg_parser);
+    return test_result ? 0 : -1;
+}

include/ck_tile/core/arch/amd_buffer_addressing_builtins.hpp

@@ -43,6 +43,11 @@ __device__ inline int32_t amd_wave_read_first_lane(int32_t value)
     return __builtin_amdgcn_readfirstlane(value);
 }
 
+__device__ inline int32_t amd_wave_read_first_lane(uintptr_t value)
+{
+    return __builtin_amdgcn_readfirstlane(value);
+}
+
 template <typename Object, std::enable_if_t<std::is_trivially_copyable_v<Object>, int> = 0>
 __device__ inline auto amd_wave_read_first_lane(const Object& obj)
 {

include/ck_tile/ops/sparse_attn.hpp

@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/ops/sparse_attn/block_fmha_pipeline_qr_ks_vs_async_jenga.hpp"
+#include "ck_tile/ops/sparse_attn/block_fmha_pipeline_qr_ks_vs_async_vsa.hpp"
+#include "ck_tile/ops/sparse_attn/fmha_fwd_jenga_kernel.hpp"
+#include "ck_tile/ops/sparse_attn/fmha_fwd_vsa_kernel.hpp"
+#include "ck_tile/ops/common/generic_2d_block_shape.hpp"
+#include "ck_tile/ops/common/load_interleaved_pk_type.hpp"
+#include "ck_tile/ops/common/streamk_common.hpp"
+#include "ck_tile/ops/common/tensor_layout.hpp"
+#include "ck_tile/ops/common/utils.hpp"

include/ck_tile/ops/sparse_attn/block_fmha_pipeline_qr_ks_vs_async_jenga.hpp

@@ -431,11 +431,6 @@ struct BlockFmhaPipelineQRKSVSAsyncJenga
                                           sequence<(LdsSeq.at(number<i_k0>{})) * kN0, 0>{},
                                           sequence<(LdsSeq.at(number<i_k0>{}) + 1) * kN0, kK0>{}));
                 });
-                __shared__ int printed_flag;
-                if(blockIdx.x == 0 && threadIdx.x == 0 && i_total_loops == 1000)
-                {
-                    printed_flag = 100;
-                }
             }
 
             // TODO: this to fix a bug when loop smaller than 2,
@@ -704,14 +699,8 @@ struct BlockFmhaPipelineQRKSVSAsyncJenga
                     randval_dram_window);
             }
 
-            const auto p = [&]() {
-                if constexpr(std::is_same_v<PDataType, fp16_t>)
-                    return impl::cast_tile_pk_fp16_fp32<PDataType>(
-                        tile_elementwise_in(p_compute_element_func, p_compute));
-                else
-                    return cast_tile<PDataType>(
-                        tile_elementwise_in(p_compute_element_func, p_compute));
-            }();
+            const auto p =
+                cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, p_compute));
 
             // STAGE 3, KV gemm
             if constexpr(k1_loops > 1)

include/ck_tile/ops/sparse_attn/fmha_fwd_jenga_kernel.hpp

@@ -53,8 +53,9 @@ struct FmhaFwdJengaKernel
     static constexpr auto BiasEnum          = FmhaPipeline::BiasEnum;
     static constexpr bool kStoreLSE         = FmhaPipeline::kStoreLSE;
     static constexpr bool kHasDropout       = FmhaPipeline::kHasDropout;
-    static constexpr bool kDoFp8StaticQuant = FmhaPipeline::Problem::kDoFp8StaticQuant;
-    static constexpr bool kSkipMinSeqlenQ   = FmhaPipeline::Problem::kSkipMinSeqlenQ;
+    static constexpr bool kDoFp8StaticQuant =
+        (FmhaPipeline::Problem::QScaleEnum != ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE);
+    static constexpr bool kSkipMinSeqlenQ = FmhaPipeline::Problem::kSkipMinSeqlenQ;
 
     using AttentionVariant = ck_tile::remove_cvref_t<typename FmhaPipeline::AttentionVariant>;
     using FmhaMask         = ck_tile::remove_cvref_t<typename FmhaPipeline::FmhaMask>;