From f9caae2d8bae2bd2ef69fa58170829b2bd419b20 Mon Sep 17 00:00:00 2001 From: Qianfeng Zhang Date: Wed, 18 Jun 2025 15:57:04 +0000 Subject: [PATCH] Use batch dim as first grid dim by default and replace env ASSUME_LEAST_VARIED_SEQLEN by ASSUME_HIGHLY_VARIED_SEQLEN --- example/ck_tile/18_hstu_attention/CMakeLists.txt | 4 ++-- .../ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/example/ck_tile/18_hstu_attention/CMakeLists.txt b/example/ck_tile/18_hstu_attention/CMakeLists.txt index c62df256f3..5145984cb9 100644 --- a/example/ck_tile/18_hstu_attention/CMakeLists.txt +++ b/example/ck_tile/18_hstu_attention/CMakeLists.txt @@ -12,8 +12,8 @@ set(EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS) list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3) -if (DEFINED ENV{ASSUME_LEAST_VARIED_SEQLEN}) - list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DHSTU_SCHED_BATCH_AS_FIRST_GRID_DIM=1) +if (DEFINED ENV{ASSUME_HIGHLY_VARIED_SEQLEN}) + list(APPEND EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS -DHSTU_SCHED_BATCH_AS_FIRST_GRID_DIM=0) endif() target_compile_options(${EXAMPLE_HSTU_ATTENTION} PRIVATE ${EXAMPLE_HSTU_ATTENTION_COMPILE_OPTIONS}) diff --git a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp index 1f1f032981..94a29a1bb4 100644 --- a/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp +++ b/example/ck_tile/18_hstu_attention/hstu_attention_fwd_kernel.hpp @@ -14,7 +14,7 @@ #include "hstu_block_masking.hpp" #ifndef HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM -#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 0 +#define HSTU_SCHED_BATCH_AS_FIRST_GRID_DIM 1 #endif // S[seqlen_q, seqlen_k] = Q[seqlen_q, hdim_q] @ K[seqlen_k, hdim_q]