diff --git a/.github/workflows/build-rocm-wheels.yml b/.github/workflows/build-rocm-wheels.yml index c7075c7e897b..f98135529a53 100644 --- a/.github/workflows/build-rocm-wheels.yml +++ b/.github/workflows/build-rocm-wheels.yml @@ -121,7 +121,13 @@ jobs: env: PYTORCH_ROCM_ARCH: ${{ env.PYTORCH_ROCM_ARCH }} VLLM_TARGET_DEVICE: rocm - MAX_JOBS: 1 + # MAX_JOBS=2 was reverted (adf8d91564) due to OOM on 16 GB runners. + # The heavy TUs have since been split into per-N shards (see + # csrc/rocm/skinny_gemms_int4/, skinny_gemms/, skinny_gemms_int8/, + # attention/, and csrc/libtorch_stable/attention/paged_attention_v1/). + # Peak RSS per shard is now <=3.2 GB (down from 12.6 GB for int4). + # Measured worst-case concurrent peak at MAX_JOBS=4: ~13 GB. + MAX_JOBS: 4 SCCACHE_GHA_ENABLED: "true" SCCACHE_IDLE_TIMEOUT: 0 run: | diff --git a/CMakeLists.txt b/CMakeLists.txt index f6bdd5770639..a4215e27ffd7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -644,7 +644,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/sampler.cu" "csrc/libtorch_stable/topk.cu" "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" + "csrc/libtorch_stable/mamba/selective_scan_fwd/shard_fp16.cu" + "csrc/libtorch_stable/mamba/selective_scan_fwd/shard_bf16.cu" "csrc/libtorch_stable/attention/paged_attention_v1.cu" + "csrc/libtorch_stable/attention/paged_attention_v1/shard_small.cu" + "csrc/libtorch_stable/attention/paged_attention_v1/shard_medium.cu" + "csrc/libtorch_stable/attention/paged_attention_v1/shard_large.cu" "csrc/libtorch_stable/attention/paged_attention_v2.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels.cu" @@ -1295,15 +1300,32 @@ if(VLLM_GPU_LANG STREQUAL "HIP") set(VLLM_ROCM_EXT_SRC "csrc/rocm/torch_bindings.cpp" "csrc/rocm/skinny_gemms.cu" + "csrc/rocm/skinny_gemms/instantiate_n1.cu" + "csrc/rocm/skinny_gemms/instantiate_n2.cu" + "csrc/rocm/skinny_gemms/instantiate_n3.cu" + "csrc/rocm/skinny_gemms/instantiate_n4.cu" + "csrc/rocm/skinny_gemms/instantiate_n5.cu" "csrc/rocm/skinny_gemms_int8.cu" + "csrc/rocm/skinny_gemms_int8/instantiate_n1.cu" + "csrc/rocm/skinny_gemms_int8/instantiate_n2.cu" + "csrc/rocm/skinny_gemms_int8/instantiate_n3.cu" + "csrc/rocm/skinny_gemms_int8/instantiate_n4.cu" + "csrc/rocm/skinny_gemms_int8/instantiate_n5.cu" "csrc/rocm/skinny_gemms_int4.cu" + "csrc/rocm/skinny_gemms_int4/instantiate_n1.cu" + "csrc/rocm/skinny_gemms_int4/instantiate_n2.cu" + "csrc/rocm/skinny_gemms_int4/instantiate_n3.cu" + "csrc/rocm/skinny_gemms_int4/instantiate_n4.cu" + "csrc/rocm/skinny_gemms_int4/instantiate_n5.cu" "csrc/rocm/skinny_gemms_w8a8.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n1.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n2.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n3.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n4.cu" "csrc/rocm/skinny_gemms_w8a8/instantiate_n5.cu" - "csrc/rocm/attention.cu") + "csrc/rocm/attention.cu" + "csrc/rocm/attention/dispatch_auto.cu" + "csrc/rocm/attention/dispatch_fp8.cu") set(VLLM_ROCM_FLAGS ${VLLM_GPU_FLAGS}) diff --git a/csrc/libtorch_stable/attention/paged_attention_v1.cu b/csrc/libtorch_stable/attention/paged_attention_v1.cu index 8fa417915936..c341b931add1 100644 --- a/csrc/libtorch_stable/attention/paged_attention_v1.cu +++ b/csrc/libtorch_stable/attention/paged_attention_v1.cu @@ -16,175 +16,61 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "../torch_utils.h" -#include "attention_kernels.cuh" -#include "../../cuda_compat.h" -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) +#include "paged_attention_v1/common.h" -#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE) \ - VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( \ - ((void*)vllm::paged_attention_v1_kernel), \ - shared_mem_size); \ - vllm::paged_attention_v1_kernel \ - <<>>( \ - out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \ - scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \ - alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \ - k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks, \ - blocksparse_vert_stride, blocksparse_block_size, \ - blocksparse_head_sliding_step); - -// TODO(woosuk): Tune NUM_THREADS. -template -void paged_attention_v1_launcher( - torch::stable::Tensor& out, torch::stable::Tensor& query, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - int num_kv_heads, float scale, torch::stable::Tensor& block_tables, - torch::stable::Tensor& seq_lens, int max_seq_len, +void paged_attention_v1( + torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] + torch::stable::Tensor& + key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] + torch::stable::Tensor& + value_cache, // [num_blocks, num_heads, head_size, block_size] + int64_t num_kv_heads, // [num_heads] + double scale, + torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::stable::Tensor& seq_lens, // [num_seqs] + int64_t block_size, int64_t max_seq_len, const std::optional& alibi_slopes, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, - const int tp_rank, const int blocksparse_local_blocks, - const int blocksparse_vert_stride, const int blocksparse_block_size, - const int blocksparse_head_sliding_step) { - int num_seqs = query.size(0); - int num_heads = query.size(1); + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step) { int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; - - T* out_ptr = reinterpret_cast(out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.mutable_data_ptr(); - int* seq_lens_ptr = seq_lens.mutable_data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - - const int NUM_WARPS = NUM_THREADS / WARP_SIZE; - int padded_max_seq_len = - DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE; - int logits_size = padded_max_seq_len * sizeof(float); - int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); - // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len - // Keep that in sync with the logic here! - int shared_mem_size = std::max(logits_size, outputs_size); - - dim3 grid(num_heads, num_seqs, 1); - dim3 block(NUM_THREADS); - const torch::stable::accelerator::DeviceGuard device_guard( - query.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); switch (head_size) { - // NOTE(woosuk): To reduce the compilation time, we only compile for the - // head sizes that we use in the model. However, we can easily extend this - // to support any head size which is a multiple of 16. case 32: - LAUNCH_PAGED_ATTENTION_V1(32); - break; case 64: - LAUNCH_PAGED_ATTENTION_V1(64); - break; case 80: - LAUNCH_PAGED_ATTENTION_V1(80); + paged_attention_v1_small( + out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, + seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, + k_scale, v_scale, tp_rank, blocksparse_local_blocks, + blocksparse_vert_stride, blocksparse_block_size, + blocksparse_head_sliding_step); break; case 96: - LAUNCH_PAGED_ATTENTION_V1(96); - break; case 112: - LAUNCH_PAGED_ATTENTION_V1(112); - break; case 120: - LAUNCH_PAGED_ATTENTION_V1(120); - break; case 128: - LAUNCH_PAGED_ATTENTION_V1(128); + paged_attention_v1_medium( + out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, + seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, + k_scale, v_scale, tp_rank, blocksparse_local_blocks, + blocksparse_vert_stride, blocksparse_block_size, + blocksparse_head_sliding_step); break; case 192: - LAUNCH_PAGED_ATTENTION_V1(192); - break; case 256: - LAUNCH_PAGED_ATTENTION_V1(256); + paged_attention_v1_large( + out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, + seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, + k_scale, v_scale, tp_rank, blocksparse_local_blocks, + blocksparse_vert_stride, blocksparse_block_size, + blocksparse_head_sliding_step); break; default: STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); break; } } - -#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ - paged_attention_v1_launcher( \ - out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \ - seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank, \ - blocksparse_local_blocks, blocksparse_vert_stride, \ - blocksparse_block_size, blocksparse_head_sliding_step); - -#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ - if (is_block_sparse) { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ - } else { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ - } - -// NOTE(woosuk): To reduce the compilation time, we omitted block sizes -// 1, 2, 4, 64, 128, 256. -#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -void paged_attention_v1( - torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, // [num_heads] - double scale, - torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::stable::Tensor& seq_lens, // [num_seqs] - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step) { - const bool is_block_sparse = (blocksparse_vert_stride > 1); - - DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, - CALL_V1_LAUNCHER_BLOCK_SIZE) -} - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/attention/paged_attention_v1/common.h b/csrc/libtorch_stable/attention/paged_attention_v1/common.h new file mode 100644 index 000000000000..86989b1c6f18 --- /dev/null +++ b/csrc/libtorch_stable/attention/paged_attention_v1/common.h @@ -0,0 +1,152 @@ +#pragma once + +#include "../../torch_utils.h" +#include "../attention_kernels.cuh" +#include "../../../cuda_compat.h" +#include "../../../quantization/w8a8/fp8/amd/quant_utils.cuh" + +#include + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) +#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) + +#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE) \ + VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( \ + ((void*)vllm::paged_attention_v1_kernel), \ + shared_mem_size); \ + vllm::paged_attention_v1_kernel \ + <<>>( \ + out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \ + scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \ + alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \ + k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks, \ + blocksparse_vert_stride, blocksparse_block_size, \ + blocksparse_head_sliding_step); + +// Launcher preamble: extract pointers and compute grid/block dims. +// Used by each shard's launcher template before its HEAD_SIZE switch. +#define V1_LAUNCHER_PREAMBLE \ + int num_seqs = query.size(0); \ + int num_heads = query.size(1); \ + int head_size = query.size(2); \ + int max_num_blocks_per_seq = block_tables.size(1); \ + int q_stride = query.stride(0); \ + int kv_block_stride = key_cache.stride(0); \ + int kv_head_stride = key_cache.stride(1); \ + \ + const float* alibi_slopes_ptr = \ + alibi_slopes \ + ? reinterpret_cast(alibi_slopes.value().data_ptr()) \ + : nullptr; \ + \ + T* out_ptr = reinterpret_cast(out.data_ptr()); \ + T* query_ptr = reinterpret_cast(query.data_ptr()); \ + CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); \ + CACHE_T* value_cache_ptr = \ + reinterpret_cast(value_cache.data_ptr()); \ + int* block_tables_ptr = block_tables.mutable_data_ptr(); \ + int* seq_lens_ptr = seq_lens.mutable_data_ptr(); \ + const float* k_scale_ptr = \ + reinterpret_cast(k_scale.data_ptr()); \ + const float* v_scale_ptr = \ + reinterpret_cast(v_scale.data_ptr()); \ + \ + const int NUM_WARPS = NUM_THREADS / WARP_SIZE; \ + int padded_max_seq_len = \ + DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE; \ + int logits_size = padded_max_seq_len * sizeof(float); \ + int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); \ + int shared_mem_size = std::max(logits_size, outputs_size); \ + \ + dim3 grid(num_heads, num_seqs, 1); \ + dim3 block(NUM_THREADS); \ + const torch::stable::accelerator::DeviceGuard device_guard( \ + query.get_device_index()); \ + const cudaStream_t stream = get_current_cuda_stream(); + +// Launcher function signature (same for all shards). +#define V1_LAUNCHER_PARAMS \ + torch::stable::Tensor &out, torch::stable::Tensor &query, \ + torch::stable::Tensor &key_cache, torch::stable::Tensor &value_cache, \ + int num_kv_heads, float scale, torch::stable::Tensor &block_tables, \ + torch::stable::Tensor &seq_lens, int max_seq_len, \ + const std::optional&alibi_slopes, \ + torch::stable::Tensor &k_scale, torch::stable::Tensor &v_scale, \ + const int tp_rank, const int blocksparse_local_blocks, \ + const int blocksparse_vert_stride, const int blocksparse_block_size, \ + const int blocksparse_head_sliding_step + +#define V1_LAUNCHER_ARGS \ + out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \ + seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank, \ + blocksparse_local_blocks, blocksparse_vert_stride, \ + blocksparse_block_size, blocksparse_head_sliding_step + +// Dispatch macros used by the entry point (same as original). +#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ + paged_attention_v1_launcher(V1_LAUNCHER_ARGS); + +#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ + if (is_block_sparse) { \ + CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ + } else { \ + CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ + } + +#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ + switch (block_size) { \ + case 8: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ + break; \ + case 16: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ + break; \ + case 32: \ + CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ + break; \ + default: \ + STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ + } + +// Per-shard entry point declarations (non-template, called from entry point). +void paged_attention_v1_small( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +void paged_attention_v1_medium( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); + +void paged_attention_v1_large( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step); diff --git a/csrc/libtorch_stable/attention/paged_attention_v1/shard_large.cu b/csrc/libtorch_stable/attention/paged_attention_v1/shard_large.cu new file mode 100644 index 000000000000..2749d0ea46d7 --- /dev/null +++ b/csrc/libtorch_stable/attention/paged_attention_v1/shard_large.cu @@ -0,0 +1,39 @@ +// HEAD_SIZE shard: 192, 256 + +#include "common.h" + +template +void paged_attention_v1_launcher(V1_LAUNCHER_PARAMS) { + V1_LAUNCHER_PREAMBLE + + switch (head_size) { + case 192: + LAUNCH_PAGED_ATTENTION_V1(192); + break; + case 256: + LAUNCH_PAGED_ATTENTION_V1(256); + break; + default: + STD_TORCH_CHECK(false, + "Unsupported head size in large shard: ", head_size); + break; + } +} + +void paged_attention_v1_large( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step) { + const bool is_block_sparse = (blocksparse_vert_stride > 1); + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, + CALL_V1_LAUNCHER_BLOCK_SIZE) +} diff --git a/csrc/libtorch_stable/attention/paged_attention_v1/shard_medium.cu b/csrc/libtorch_stable/attention/paged_attention_v1/shard_medium.cu new file mode 100644 index 000000000000..ebebe3da4298 --- /dev/null +++ b/csrc/libtorch_stable/attention/paged_attention_v1/shard_medium.cu @@ -0,0 +1,45 @@ +// HEAD_SIZE shard: 96, 112, 120, 128 + +#include "common.h" + +template +void paged_attention_v1_launcher(V1_LAUNCHER_PARAMS) { + V1_LAUNCHER_PREAMBLE + + switch (head_size) { + case 96: + LAUNCH_PAGED_ATTENTION_V1(96); + break; + case 112: + LAUNCH_PAGED_ATTENTION_V1(112); + break; + case 120: + LAUNCH_PAGED_ATTENTION_V1(120); + break; + case 128: + LAUNCH_PAGED_ATTENTION_V1(128); + break; + default: + STD_TORCH_CHECK(false, + "Unsupported head size in medium shard: ", head_size); + break; + } +} + +void paged_attention_v1_medium( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step) { + const bool is_block_sparse = (blocksparse_vert_stride > 1); + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, + CALL_V1_LAUNCHER_BLOCK_SIZE) +} diff --git a/csrc/libtorch_stable/attention/paged_attention_v1/shard_small.cu b/csrc/libtorch_stable/attention/paged_attention_v1/shard_small.cu new file mode 100644 index 000000000000..febf79496059 --- /dev/null +++ b/csrc/libtorch_stable/attention/paged_attention_v1/shard_small.cu @@ -0,0 +1,42 @@ +// HEAD_SIZE shard: 32, 64, 80 + +#include "common.h" + +template +void paged_attention_v1_launcher(V1_LAUNCHER_PARAMS) { + V1_LAUNCHER_PREAMBLE + + switch (head_size) { + case 32: + LAUNCH_PAGED_ATTENTION_V1(32); + break; + case 64: + LAUNCH_PAGED_ATTENTION_V1(64); + break; + case 80: + LAUNCH_PAGED_ATTENTION_V1(80); + break; + default: + STD_TORCH_CHECK(false, + "Unsupported head size in small shard: ", head_size); + break; + } +} + +void paged_attention_v1_small( + torch::stable::Tensor& out, torch::stable::Tensor& query, + torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, + int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, + torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, + torch::stable::Tensor& v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step) { + const bool is_block_sparse = (blocksparse_vert_stride > 1); + DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, + CALL_V1_LAUNCHER_BLOCK_SIZE) +} diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu b/csrc/libtorch_stable/mamba/selective_scan_fwd.cu index c17b06a79bd0..ce681722f585 100644 --- a/csrc/libtorch_stable/mamba/selective_scan_fwd.cu +++ b/csrc/libtorch_stable/mamba/selective_scan_fwd.cu @@ -4,460 +4,7 @@ #include #include "selective_scan.h" -#ifndef USE_ROCM - #include - #include - #include -#else - #include - namespace cub = hipcub; -#endif - -#include "selective_scan.h" -#include "static_switch.h" - -template -struct Selective_Scan_fwd_kernel_traits { - static_assert(kNItems_ % 4 == 0); - using input_t = input_t_; - using weight_t = weight_t_; - using state_t = state_t_; - static constexpr int kNThreads = kNThreads_; - // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy. - static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3; - static constexpr int kNItems = kNItems_; - static constexpr int kNRows = kNRows_; - static constexpr int kNBytes = sizeof(input_t); - static_assert(kNBytes == 2 || kNBytes == 4); - static constexpr int kNElts = kNBytes == 4 ? 4 : constexpr_min(8, kNItems); - static_assert(kNItems % kNElts == 0); - static constexpr int kNLoads = kNItems / kNElts; - static constexpr bool kIsEvenLen = kVarlen_ ? false : kIsEvenLen_; - static constexpr bool kIsVariableB = kIsVariableB_; - static constexpr bool kIsVariableC = kIsVariableC_; - static constexpr bool kHasZ = kHasZ_; - static constexpr bool kVarlen = kVarlen_; - - static constexpr bool kDirectIO = kVarlen_ ? false : kIsEvenLen && kNLoads == 1; - static constexpr int kNLoadsIndex = kNItems / 4; - using vec_t = typename BytesToType::Type; - using scan_t = float2; - using BlockLoadT = cub::BlockLoad; - using BlockLoadVecT = cub::BlockLoad; - using BlockLoadWeightT = cub::BlockLoad; - using BlockLoadWeightVecT = cub::BlockLoad; - using BlockStoreT = cub::BlockStore; - using BlockStoreVecT = cub::BlockStore; - // using BlockScanT = cub::BlockScan; - // using BlockScanT = cub::BlockScan; - using BlockScanT = cub::BlockScan; - static constexpr int kSmemIOSize = custom_max({sizeof(typename BlockLoadT::TempStorage), - sizeof(typename BlockLoadVecT::TempStorage), - (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightT::TempStorage), - (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightVecT::TempStorage), - sizeof(typename BlockStoreT::TempStorage), - sizeof(typename BlockStoreVecT::TempStorage)}); - static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage); -}; - -template -__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks) -void selective_scan_fwd_kernel(SSMParamsBase params) { - constexpr bool kIsVariableB = Ktraits::kIsVariableB; - constexpr bool kIsVariableC = Ktraits::kIsVariableC; - constexpr bool kHasZ = Ktraits::kHasZ; - constexpr bool kVarlen = Ktraits::kVarlen; - constexpr int kNItems = Ktraits::kNItems; - constexpr int kNRows = Ktraits::kNRows; - constexpr bool kDirectIO = Ktraits::kDirectIO; - using input_t = typename Ktraits::input_t; - using weight_t = typename Ktraits::weight_t; - using scan_t = typename Ktraits::scan_t; - - // Shared memory. - extern __shared__ char smem_[]; - // cast to lvalue reference of expected type - // char *smem_loadstorescan = smem_ + 2 * MAX_DSTATE * sizeof(weight_t); - // auto& smem_load = reinterpret_cast(smem_ + 2 * MAX_DSTATE * sizeof(weight_t)); - // auto& smem_load = reinterpret_cast(smem_loadstorescan); - auto& smem_load = reinterpret_cast(smem_); - auto& smem_load_weight = reinterpret_cast(smem_); - auto& smem_load_weight1 = *reinterpret_cast(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage)); - auto& smem_store = reinterpret_cast(smem_); - auto& smem_scan = *reinterpret_cast(smem_ + Ktraits::kSmemIOSize); - // weight_t *smem_a = reinterpret_cast(smem_ + smem_loadstorescan_size); - // weight_t *smem_bc = reinterpret_cast(smem_a + MAX_DSTATE); - scan_t *smem_running_prefix = reinterpret_cast(smem_ + Ktraits::kSmemSize); - - const int batch_id = blockIdx.x; - const int dim_id = blockIdx.y; - const int group_id = dim_id / (params.dim_ngroups_ratio); - int seqlen = params.seqlen; - int sequence_start_index = batch_id; - if constexpr (kVarlen){ - int *query_start_loc = reinterpret_cast(params.query_start_loc_ptr); - sequence_start_index = query_start_loc[batch_id]; - seqlen = query_start_loc[batch_id + 1] - sequence_start_index; - } - const bool has_initial_state = params.has_initial_state_ptr == nullptr ? false - : reinterpret_cast(params.has_initial_state_ptr)[batch_id]; - - const int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr - : reinterpret_cast(params.cache_indices_ptr); - int cache_index; - if (cache_indices == nullptr) { - cache_index = batch_id; - } else if (params.cache_enabled) { - const int* initial_state_idx = reinterpret_cast(params.initial_state_idx_ptr); - cache_index = cache_indices[batch_id * params.cache_indices_stride + initial_state_idx[batch_id]]; - } else { - cache_index = cache_indices[batch_id]; - } - // Skip batch entries whose cache index maps to the null block (padding). - if (cache_indices != nullptr && cache_index == params.null_block_id){ - return; - } - input_t *u = reinterpret_cast(params.u_ptr) + sequence_start_index * params.u_batch_stride - + dim_id * kNRows * params.u_d_stride; - input_t *delta = reinterpret_cast(params.delta_ptr) + sequence_start_index * params.delta_batch_stride - + dim_id * kNRows * params.delta_d_stride; - weight_t *A = reinterpret_cast(params.A_ptr) + dim_id * kNRows * params.A_d_stride; - weight_t *B = reinterpret_cast(params.B_ptr) + dim_id * kNRows * params.B_d_stride; - input_t *Bvar = reinterpret_cast(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride; - weight_t *C = reinterpret_cast(params.C_ptr) + dim_id * kNRows * params.C_d_stride; - input_t *Cvar = reinterpret_cast(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride; - - typename Ktraits::state_t *ssm_states; - if (params.cache_enabled) { - // APC mode: ssm_states points to the base, we'll use absolute cache slots later - ssm_states = reinterpret_cast(params.ssm_states_ptr) + - dim_id * kNRows * params.ssm_states_dim_stride; - } else { - // Non-APC mode: offset by cache_index as before - ssm_states = reinterpret_cast(params.ssm_states_ptr) + - cache_index * params.ssm_states_batch_stride + - dim_id * kNRows * params.ssm_states_dim_stride; - } - - float D_val[kNRows] = {0}; - if (params.D_ptr != nullptr) { - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - D_val[r] = reinterpret_cast(params.D_ptr)[dim_id * kNRows + r]; - } - } - float delta_bias[kNRows] = {0}; - if (params.delta_bias_ptr != nullptr) { - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - delta_bias[r] = reinterpret_cast(params.delta_bias_ptr)[dim_id * kNRows + r]; - } - } - - // Use block_size for chunking when APC is enabled, otherwise use 2048 for backwards compatibility - const int block_size = params.cache_enabled ? params.block_size : 2048; - - const int* batch_cache_indices = cache_indices != nullptr ? - cache_indices + batch_id * params.cache_indices_stride : nullptr; - const int* block_idx_first_scheduled = params.block_idx_first_scheduled_token_ptr != nullptr ? - reinterpret_cast(params.block_idx_first_scheduled_token_ptr) : nullptr; - const int* block_idx_last_scheduled = params.block_idx_last_scheduled_token_ptr != nullptr ? - reinterpret_cast(params.block_idx_last_scheduled_token_ptr) : nullptr; - const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ? - reinterpret_cast(params.initial_state_idx_ptr) : nullptr; - const int* cu_chunk_seqlen = params.cu_chunk_seqlen_ptr != nullptr ? - reinterpret_cast(params.cu_chunk_seqlen_ptr) : nullptr; - const int* last_chunk_indices = params.last_chunk_indices_ptr != nullptr ? - reinterpret_cast(params.last_chunk_indices_ptr) : nullptr; - - const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index; - - const int block_idx_first = (params.cache_enabled && block_idx_first_scheduled != nullptr) ? - block_idx_first_scheduled[batch_id] : 0; - - // Determine chunk boundaries from pre-computed metadata (APC mode) - // or fall back to simple block_size chunking. - int first_chunk_idx, n_chunks; - int current_position; - - if (cu_chunk_seqlen != nullptr && last_chunk_indices != nullptr) { - const int last_chunk_idx = last_chunk_indices[batch_id]; - first_chunk_idx = (batch_id == 0) ? 0 : last_chunk_indices[batch_id - 1] + 1; - n_chunks = last_chunk_idx - first_chunk_idx + 1; - // Derive current_position: if the first chunk is partial (fills remainder - // of a started block), offset into the block accordingly. - const int first_chunk_tokens = cu_chunk_seqlen[first_chunk_idx + 1] - cu_chunk_seqlen[first_chunk_idx]; - const int chunk_start_offset = (n_chunks > 1 && first_chunk_tokens < block_size) - ? (block_size - first_chunk_tokens) : 0; - current_position = block_idx_first * block_size + chunk_start_offset; - } else { - first_chunk_idx = 0; - n_chunks = (seqlen + block_size - 1) / block_size; - current_position = 0; - } - - int tokens_processed = 0; - - for (int chunk = 0; chunk < n_chunks; ++chunk) { - const int chunk_tokens = (cu_chunk_seqlen != nullptr) - ? cu_chunk_seqlen[first_chunk_idx + chunk + 1] - cu_chunk_seqlen[first_chunk_idx + chunk] - : min(block_size, seqlen - tokens_processed); - if (chunk_tokens <= 0) break; - input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems]; - - __syncthreads(); - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - if constexpr (!kDirectIO) { - if (r > 0) { __syncthreads(); } - } - load_input(u + r * params.u_d_stride, u_vals[r], smem_load, chunk_tokens); - if constexpr (!kDirectIO) { __syncthreads(); } - load_input(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, chunk_tokens); - } - u += chunk_tokens; - delta += chunk_tokens; - - float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems]; - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - #pragma unroll - for (int i = 0; i < kNItems; ++i) { - float u_val = float(u_vals[r][i]); - delta_vals[r][i] = float(delta_vals_load[r][i]) + delta_bias[r]; - if (params.delta_softplus) { - delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i]; - } - delta_u_vals[r][i] = delta_vals[r][i] * u_val; - out_vals[r][i] = D_val[r] * u_val; - } - } - - __syncthreads(); - for (int state_idx = 0; state_idx < params.dstate; ++state_idx) { - weight_t A_val[kNRows]; - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride]; - // Multiply the real part of A with LOG2E so we can use exp2f instead of expf. - constexpr float kLog2e = M_LOG2E; - A_val[r] *= kLog2e; - } - // This variable holds B * C if both B and C are constant across seqlen. If only B varies - // across seqlen, this holds C. If only C varies across seqlen, this holds B. - // If both B and C vary, this is unused. - weight_t BC_val[kNRows]; - weight_t B_vals[kNItems], C_vals[kNItems]; - if constexpr (kIsVariableB) { - load_weight(Bvar + state_idx * params.B_dstate_stride, B_vals, - smem_load_weight, chunk_tokens); - if constexpr (!kIsVariableC) { - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - BC_val[r] = C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; - } - } - } - if constexpr (kIsVariableC) { - auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1; - load_weight(Cvar + state_idx * params.C_dstate_stride, C_vals, - smem_load_weight_C, chunk_tokens); - if constexpr (!kIsVariableB) { - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride]; - } - } - } - if constexpr (!kIsVariableB && !kIsVariableC) { - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride] * C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; - } - } - - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - if (r > 0) { __syncthreads(); } // Scan could be using the same smem - scan_t thread_data[kNItems]; - #pragma unroll - for (int i = 0; i < kNItems; ++i) { - thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]), - !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]); - if (threadIdx.x * kNItems + i >= chunk_tokens) { - thread_data[i] = make_float2(1.f, 0.f); - } - } - // Initialize running total - scan_t running_prefix; - if (chunk > 0) { - running_prefix = smem_running_prefix[state_idx + r * MAX_DSTATE]; - } else { - // Load initial state - if (params.cache_enabled && has_initial_state && batch_cache_indices != nullptr) { - size_t state_offset = load_cache_slot * params.ssm_states_batch_stride + - r * params.ssm_states_dim_stride + - state_idx * params.ssm_states_dstate_stride; - running_prefix = make_float2(1.0, float(ssm_states[state_offset])); - } else if (has_initial_state) { - // Non-APC mode: load from current batch position - running_prefix = make_float2(1.0, float(ssm_states[state_idx * params.ssm_states_dstate_stride])); - } else { - // No initial state - running_prefix = make_float2(1.0, 0.0); - } - } - - SSMScanPrefixCallbackOp prefix_op(running_prefix); - typename Ktraits::BlockScanT(smem_scan).InclusiveScan( - thread_data, thread_data, SSMScanOp(), prefix_op - ); - // There's a syncthreads in the scan op, so we don't need to sync here. - // Unless there's only 1 warp, but then it's the same thread (0) reading and writing. - if (threadIdx.x == 0) { - smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix; - - // Store state at the end of each aligned chunk when cache is enabled - if (params.cache_enabled && batch_cache_indices != nullptr) { - size_t cache_slot; - if (chunk == n_chunks - 1) { - cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]]; - } else { - const int block_idx_completed = (current_position + chunk_tokens - 1) / block_size; - cache_slot = batch_cache_indices[block_idx_completed]; - } - - size_t state_offset = cache_slot * params.ssm_states_batch_stride + - r * params.ssm_states_dim_stride + - state_idx * params.ssm_states_dstate_stride; - - ssm_states[state_offset] = typename Ktraits::state_t(prefix_op.running_prefix.y); - } else if (!params.cache_enabled && chunk == n_chunks - 1) { - // Non-APC mode: store only final state at current batch position - ssm_states[state_idx * params.ssm_states_dstate_stride] = typename Ktraits::state_t(prefix_op.running_prefix.y); - } - } - #pragma unroll - for (int i = 0; i < kNItems; ++i) { - const weight_t C_val = !kIsVariableC - ? BC_val[r] - : (!kIsVariableB ? BC_val[r] * C_vals[i] : C_vals[i]); - out_vals[r][i] += thread_data[i].y * C_val; - } - } - } - input_t *out = reinterpret_cast(params.out_ptr) + sequence_start_index * params.out_batch_stride - + dim_id * kNRows * params.out_d_stride + tokens_processed; - __syncthreads(); - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - if constexpr (!kDirectIO) { - if (r > 0) { __syncthreads(); } - } - store_output(out + r * params.out_d_stride, out_vals[r], smem_store, chunk_tokens); - } - - if constexpr (kHasZ) { - input_t *z = reinterpret_cast(params.z_ptr) + sequence_start_index * params.z_batch_stride - + dim_id * kNRows * params.z_d_stride + tokens_processed; - input_t *out_z = reinterpret_cast(params.out_z_ptr) + sequence_start_index * params.out_z_batch_stride - + dim_id * kNRows * params.out_z_d_stride + tokens_processed; - #pragma unroll - for (int r = 0; r < kNRows; ++r) { - input_t z_vals[kNItems]; - __syncthreads(); - load_input(z + r * params.z_d_stride, z_vals, smem_load, chunk_tokens); - #pragma unroll - for (int i = 0; i < kNItems; ++i) { - float z_val = z_vals[i]; - out_vals[r][i] *= z_val / (1 + expf(-z_val)); - } - __syncthreads(); - store_output(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, chunk_tokens); - } - } - - Bvar += chunk_tokens; - Cvar += chunk_tokens; - - tokens_processed += chunk_tokens; - current_position += chunk_tokens; - } -} - -template -void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) { - // Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block - // processing 1 row. - constexpr int kNRows = 1; - // kIsVariableB, kIsVariableC and kHasZ are all set to True to reduce binary size - constexpr bool kIsVariableB = true; - constexpr bool kIsVariableC = true; - BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] { - BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] { - BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] { - using Ktraits = Selective_Scan_fwd_kernel_traits; - constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); - dim3 grid(params.batch, params.dim / kNRows); - auto kernel = &selective_scan_fwd_kernel; - if (kSmemSize >= 48 * 1024) { -#ifdef USE_ROCM - STD_CUDA_CHECK(hipFuncSetAttribute( - reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); -#else - STD_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); -#endif - } - kernel<<>>(params); - STD_CUDA_KERNEL_LAUNCH_CHECK(); - }); - }); - }); -} - -template -void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) { - - #ifndef USE_ROCM - if (params.cache_enabled && params.block_size == 1024) { - selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 128) { - selective_scan_fwd_launch<32, 4, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 256) { - selective_scan_fwd_launch<32, 8, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 512) { - selective_scan_fwd_launch<32, 16, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 1024) { - selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); - } else { - selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); - } - #else - if (params.cache_enabled && params.block_size == 1024) { - selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 256) { - selective_scan_fwd_launch<64, 4, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 512) { - selective_scan_fwd_launch<64, 8, input_t, weight_t, state_t>(params, stream); - } else if (params.seqlen <= 1024) { - selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); - } else { - selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); - } - #endif -} - -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); +#include #define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")") @@ -636,7 +183,7 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms, } params.out_batch_stride = out.stride(0); params.out_d_stride = out.stride(1); - + params.ssm_states_batch_stride = ssm_states.stride(0); params.ssm_states_dim_stride = ssm_states.stride(1); params.ssm_states_dstate_stride = ssm_states.stride(2); @@ -813,7 +360,7 @@ void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Ten last_chunk_indices ); - + const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index()); auto stream = get_current_cuda_stream(); DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] { diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd/kernel.cuh b/csrc/libtorch_stable/mamba/selective_scan_fwd/kernel.cuh new file mode 100644 index 000000000000..ee9796be4ce8 --- /dev/null +++ b/csrc/libtorch_stable/mamba/selective_scan_fwd/kernel.cuh @@ -0,0 +1,456 @@ +// clang-format off +// adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh +#pragma once + +#include "../../torch_utils.h" +#include +#include "../selective_scan.h" + +#ifndef USE_ROCM + #include + #include + #include +#else + #include + namespace cub = hipcub; +#endif + +#include "../selective_scan.h" +#include "../static_switch.h" + +template +struct Selective_Scan_fwd_kernel_traits { + static_assert(kNItems_ % 4 == 0); + using input_t = input_t_; + using weight_t = weight_t_; + using state_t = state_t_; + static constexpr int kNThreads = kNThreads_; + // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy. + static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3; + static constexpr int kNItems = kNItems_; + static constexpr int kNRows = kNRows_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 4 : constexpr_min(8, kNItems); + static_assert(kNItems % kNElts == 0); + static constexpr int kNLoads = kNItems / kNElts; + static constexpr bool kIsEvenLen = kVarlen_ ? false : kIsEvenLen_; + static constexpr bool kIsVariableB = kIsVariableB_; + static constexpr bool kIsVariableC = kIsVariableC_; + static constexpr bool kHasZ = kHasZ_; + static constexpr bool kVarlen = kVarlen_; + + static constexpr bool kDirectIO = kVarlen_ ? false : kIsEvenLen && kNLoads == 1; + static constexpr int kNLoadsIndex = kNItems / 4; + using vec_t = typename BytesToType::Type; + using scan_t = float2; + using BlockLoadT = cub::BlockLoad; + using BlockLoadVecT = cub::BlockLoad; + using BlockLoadWeightT = cub::BlockLoad; + using BlockLoadWeightVecT = cub::BlockLoad; + using BlockStoreT = cub::BlockStore; + using BlockStoreVecT = cub::BlockStore; + // using BlockScanT = cub::BlockScan; + // using BlockScanT = cub::BlockScan; + using BlockScanT = cub::BlockScan; + static constexpr int kSmemIOSize = custom_max({sizeof(typename BlockLoadT::TempStorage), + sizeof(typename BlockLoadVecT::TempStorage), + (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightT::TempStorage), + (int(kIsVariableB) + int(kIsVariableC)) * sizeof(typename BlockLoadWeightVecT::TempStorage), + sizeof(typename BlockStoreT::TempStorage), + sizeof(typename BlockStoreVecT::TempStorage)}); + static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage); +}; + +template +__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks) +void selective_scan_fwd_kernel(SSMParamsBase params) { + constexpr bool kIsVariableB = Ktraits::kIsVariableB; + constexpr bool kIsVariableC = Ktraits::kIsVariableC; + constexpr bool kHasZ = Ktraits::kHasZ; + constexpr bool kVarlen = Ktraits::kVarlen; + constexpr int kNItems = Ktraits::kNItems; + constexpr int kNRows = Ktraits::kNRows; + constexpr bool kDirectIO = Ktraits::kDirectIO; + using input_t = typename Ktraits::input_t; + using weight_t = typename Ktraits::weight_t; + using scan_t = typename Ktraits::scan_t; + + // Shared memory. + extern __shared__ char smem_[]; + // cast to lvalue reference of expected type + // char *smem_loadstorescan = smem_ + 2 * MAX_DSTATE * sizeof(weight_t); + // auto& smem_load = reinterpret_cast(smem_ + 2 * MAX_DSTATE * sizeof(weight_t)); + // auto& smem_load = reinterpret_cast(smem_loadstorescan); + auto& smem_load = reinterpret_cast(smem_); + auto& smem_load_weight = reinterpret_cast(smem_); + auto& smem_load_weight1 = *reinterpret_cast(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage)); + auto& smem_store = reinterpret_cast(smem_); + auto& smem_scan = *reinterpret_cast(smem_ + Ktraits::kSmemIOSize); + // weight_t *smem_a = reinterpret_cast(smem_ + smem_loadstorescan_size); + // weight_t *smem_bc = reinterpret_cast(smem_a + MAX_DSTATE); + scan_t *smem_running_prefix = reinterpret_cast(smem_ + Ktraits::kSmemSize); + + const int batch_id = blockIdx.x; + const int dim_id = blockIdx.y; + const int group_id = dim_id / (params.dim_ngroups_ratio); + int seqlen = params.seqlen; + int sequence_start_index = batch_id; + if constexpr (kVarlen){ + int *query_start_loc = reinterpret_cast(params.query_start_loc_ptr); + sequence_start_index = query_start_loc[batch_id]; + seqlen = query_start_loc[batch_id + 1] - sequence_start_index; + } + const bool has_initial_state = params.has_initial_state_ptr == nullptr ? false + : reinterpret_cast(params.has_initial_state_ptr)[batch_id]; + + const int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr + : reinterpret_cast(params.cache_indices_ptr); + int cache_index; + if (cache_indices == nullptr) { + cache_index = batch_id; + } else if (params.cache_enabled) { + const int* initial_state_idx = reinterpret_cast(params.initial_state_idx_ptr); + cache_index = cache_indices[batch_id * params.cache_indices_stride + initial_state_idx[batch_id]]; + } else { + cache_index = cache_indices[batch_id]; + } + // Skip batch entries whose cache index maps to the null block (padding). + if (cache_indices != nullptr && cache_index == params.null_block_id){ + return; + } + input_t *u = reinterpret_cast(params.u_ptr) + sequence_start_index * params.u_batch_stride + + dim_id * kNRows * params.u_d_stride; + input_t *delta = reinterpret_cast(params.delta_ptr) + sequence_start_index * params.delta_batch_stride + + dim_id * kNRows * params.delta_d_stride; + weight_t *A = reinterpret_cast(params.A_ptr) + dim_id * kNRows * params.A_d_stride; + weight_t *B = reinterpret_cast(params.B_ptr) + dim_id * kNRows * params.B_d_stride; + input_t *Bvar = reinterpret_cast(params.B_ptr) + sequence_start_index * params.B_batch_stride + group_id * params.B_group_stride; + weight_t *C = reinterpret_cast(params.C_ptr) + dim_id * kNRows * params.C_d_stride; + input_t *Cvar = reinterpret_cast(params.C_ptr) + sequence_start_index * params.C_batch_stride + group_id * params.C_group_stride; + + typename Ktraits::state_t *ssm_states; + if (params.cache_enabled) { + // APC mode: ssm_states points to the base, we'll use absolute cache slots later + ssm_states = reinterpret_cast(params.ssm_states_ptr) + + dim_id * kNRows * params.ssm_states_dim_stride; + } else { + // Non-APC mode: offset by cache_index as before + ssm_states = reinterpret_cast(params.ssm_states_ptr) + + cache_index * params.ssm_states_batch_stride + + dim_id * kNRows * params.ssm_states_dim_stride; + } + + float D_val[kNRows] = {0}; + if (params.D_ptr != nullptr) { + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + D_val[r] = reinterpret_cast(params.D_ptr)[dim_id * kNRows + r]; + } + } + float delta_bias[kNRows] = {0}; + if (params.delta_bias_ptr != nullptr) { + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + delta_bias[r] = reinterpret_cast(params.delta_bias_ptr)[dim_id * kNRows + r]; + } + } + + // Use block_size for chunking when APC is enabled, otherwise use 2048 for backwards compatibility + const int block_size = params.cache_enabled ? params.block_size : 2048; + + const int* batch_cache_indices = cache_indices != nullptr ? + cache_indices + batch_id * params.cache_indices_stride : nullptr; + const int* block_idx_first_scheduled = params.block_idx_first_scheduled_token_ptr != nullptr ? + reinterpret_cast(params.block_idx_first_scheduled_token_ptr) : nullptr; + const int* block_idx_last_scheduled = params.block_idx_last_scheduled_token_ptr != nullptr ? + reinterpret_cast(params.block_idx_last_scheduled_token_ptr) : nullptr; + const int* initial_state_idx = params.initial_state_idx_ptr != nullptr ? + reinterpret_cast(params.initial_state_idx_ptr) : nullptr; + const int* cu_chunk_seqlen = params.cu_chunk_seqlen_ptr != nullptr ? + reinterpret_cast(params.cu_chunk_seqlen_ptr) : nullptr; + const int* last_chunk_indices = params.last_chunk_indices_ptr != nullptr ? + reinterpret_cast(params.last_chunk_indices_ptr) : nullptr; + + const size_t load_cache_slot = params.cache_enabled && batch_cache_indices != nullptr ? batch_cache_indices[initial_state_idx[batch_id]] : cache_index; + + const int block_idx_first = (params.cache_enabled && block_idx_first_scheduled != nullptr) ? + block_idx_first_scheduled[batch_id] : 0; + + // Determine chunk boundaries from pre-computed metadata (APC mode) + // or fall back to simple block_size chunking. + int first_chunk_idx, n_chunks; + int current_position; + + if (cu_chunk_seqlen != nullptr && last_chunk_indices != nullptr) { + const int last_chunk_idx = last_chunk_indices[batch_id]; + first_chunk_idx = (batch_id == 0) ? 0 : last_chunk_indices[batch_id - 1] + 1; + n_chunks = last_chunk_idx - first_chunk_idx + 1; + // Derive current_position: if the first chunk is partial (fills remainder + // of a started block), offset into the block accordingly. + const int first_chunk_tokens = cu_chunk_seqlen[first_chunk_idx + 1] - cu_chunk_seqlen[first_chunk_idx]; + const int chunk_start_offset = (n_chunks > 1 && first_chunk_tokens < block_size) + ? (block_size - first_chunk_tokens) : 0; + current_position = block_idx_first * block_size + chunk_start_offset; + } else { + first_chunk_idx = 0; + n_chunks = (seqlen + block_size - 1) / block_size; + current_position = 0; + } + + int tokens_processed = 0; + + for (int chunk = 0; chunk < n_chunks; ++chunk) { + const int chunk_tokens = (cu_chunk_seqlen != nullptr) + ? cu_chunk_seqlen[first_chunk_idx + chunk + 1] - cu_chunk_seqlen[first_chunk_idx + chunk] + : min(block_size, seqlen - tokens_processed); + if (chunk_tokens <= 0) break; + input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems]; + + __syncthreads(); + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + if constexpr (!kDirectIO) { + if (r > 0) { __syncthreads(); } + } + load_input(u + r * params.u_d_stride, u_vals[r], smem_load, chunk_tokens); + if constexpr (!kDirectIO) { __syncthreads(); } + load_input(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, chunk_tokens); + } + u += chunk_tokens; + delta += chunk_tokens; + + float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems]; + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + #pragma unroll + for (int i = 0; i < kNItems; ++i) { + float u_val = float(u_vals[r][i]); + delta_vals[r][i] = float(delta_vals_load[r][i]) + delta_bias[r]; + if (params.delta_softplus) { + delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i]; + } + delta_u_vals[r][i] = delta_vals[r][i] * u_val; + out_vals[r][i] = D_val[r] * u_val; + } + } + + __syncthreads(); + for (int state_idx = 0; state_idx < params.dstate; ++state_idx) { + weight_t A_val[kNRows]; + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride]; + // Multiply the real part of A with LOG2E so we can use exp2f instead of expf. + constexpr float kLog2e = M_LOG2E; + A_val[r] *= kLog2e; + } + // This variable holds B * C if both B and C are constant across seqlen. If only B varies + // across seqlen, this holds C. If only C varies across seqlen, this holds B. + // If both B and C vary, this is unused. + weight_t BC_val[kNRows]; + weight_t B_vals[kNItems], C_vals[kNItems]; + if constexpr (kIsVariableB) { + load_weight(Bvar + state_idx * params.B_dstate_stride, B_vals, + smem_load_weight, chunk_tokens); + if constexpr (!kIsVariableC) { + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + BC_val[r] = C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; + } + } + } + if constexpr (kIsVariableC) { + auto &smem_load_weight_C = !kIsVariableB ? smem_load_weight : smem_load_weight1; + load_weight(Cvar + state_idx * params.C_dstate_stride, C_vals, + smem_load_weight_C, chunk_tokens); + if constexpr (!kIsVariableB) { + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride]; + } + } + } + if constexpr (!kIsVariableB && !kIsVariableC) { + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + BC_val[r] = B[state_idx * params.B_dstate_stride + r * params.B_d_stride] * C[state_idx * params.C_dstate_stride + r * params.C_d_stride]; + } + } + + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + if (r > 0) { __syncthreads(); } // Scan could be using the same smem + scan_t thread_data[kNItems]; + #pragma unroll + for (int i = 0; i < kNItems; ++i) { + thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]), + !kIsVariableB ? delta_u_vals[r][i] : B_vals[i] * delta_u_vals[r][i]); + if (threadIdx.x * kNItems + i >= chunk_tokens) { + thread_data[i] = make_float2(1.f, 0.f); + } + } + // Initialize running total + scan_t running_prefix; + if (chunk > 0) { + running_prefix = smem_running_prefix[state_idx + r * MAX_DSTATE]; + } else { + // Load initial state + if (params.cache_enabled && has_initial_state && batch_cache_indices != nullptr) { + size_t state_offset = load_cache_slot * params.ssm_states_batch_stride + + r * params.ssm_states_dim_stride + + state_idx * params.ssm_states_dstate_stride; + running_prefix = make_float2(1.0, float(ssm_states[state_offset])); + } else if (has_initial_state) { + // Non-APC mode: load from current batch position + running_prefix = make_float2(1.0, float(ssm_states[state_idx * params.ssm_states_dstate_stride])); + } else { + // No initial state + running_prefix = make_float2(1.0, 0.0); + } + } + + SSMScanPrefixCallbackOp prefix_op(running_prefix); + typename Ktraits::BlockScanT(smem_scan).InclusiveScan( + thread_data, thread_data, SSMScanOp(), prefix_op + ); + // There's a syncthreads in the scan op, so we don't need to sync here. + // Unless there's only 1 warp, but then it's the same thread (0) reading and writing. + if (threadIdx.x == 0) { + smem_running_prefix[state_idx + r * MAX_DSTATE] = prefix_op.running_prefix; + + // Store state at the end of each aligned chunk when cache is enabled + if (params.cache_enabled && batch_cache_indices != nullptr) { + size_t cache_slot; + if (chunk == n_chunks - 1) { + cache_slot = batch_cache_indices[block_idx_last_scheduled[batch_id]]; + } else { + const int block_idx_completed = (current_position + chunk_tokens - 1) / block_size; + cache_slot = batch_cache_indices[block_idx_completed]; + } + + size_t state_offset = cache_slot * params.ssm_states_batch_stride + + r * params.ssm_states_dim_stride + + state_idx * params.ssm_states_dstate_stride; + + ssm_states[state_offset] = typename Ktraits::state_t(prefix_op.running_prefix.y); + } else if (!params.cache_enabled && chunk == n_chunks - 1) { + // Non-APC mode: store only final state at current batch position + ssm_states[state_idx * params.ssm_states_dstate_stride] = typename Ktraits::state_t(prefix_op.running_prefix.y); + } + } + #pragma unroll + for (int i = 0; i < kNItems; ++i) { + const weight_t C_val = !kIsVariableC + ? BC_val[r] + : (!kIsVariableB ? BC_val[r] * C_vals[i] : C_vals[i]); + out_vals[r][i] += thread_data[i].y * C_val; + } + } + } + input_t *out = reinterpret_cast(params.out_ptr) + sequence_start_index * params.out_batch_stride + + dim_id * kNRows * params.out_d_stride + tokens_processed; + __syncthreads(); + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + if constexpr (!kDirectIO) { + if (r > 0) { __syncthreads(); } + } + store_output(out + r * params.out_d_stride, out_vals[r], smem_store, chunk_tokens); + } + + if constexpr (kHasZ) { + input_t *z = reinterpret_cast(params.z_ptr) + sequence_start_index * params.z_batch_stride + + dim_id * kNRows * params.z_d_stride + tokens_processed; + input_t *out_z = reinterpret_cast(params.out_z_ptr) + sequence_start_index * params.out_z_batch_stride + + dim_id * kNRows * params.out_z_d_stride + tokens_processed; + #pragma unroll + for (int r = 0; r < kNRows; ++r) { + input_t z_vals[kNItems]; + __syncthreads(); + load_input(z + r * params.z_d_stride, z_vals, smem_load, chunk_tokens); + #pragma unroll + for (int i = 0; i < kNItems; ++i) { + float z_val = z_vals[i]; + out_vals[r][i] *= z_val / (1 + expf(-z_val)); + } + __syncthreads(); + store_output(out_z + r * params.out_z_d_stride, out_vals[r], smem_store, chunk_tokens); + } + } + + Bvar += chunk_tokens; + Cvar += chunk_tokens; + + tokens_processed += chunk_tokens; + current_position += chunk_tokens; + } +} + +template +void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) { + // Only kNRows == 1 is tested for now, which ofc doesn't differ from previously when we had each block + // processing 1 row. + constexpr int kNRows = 1; + // kIsVariableB, kIsVariableC and kHasZ are all set to True to reduce binary size + constexpr bool kIsVariableB = true; + constexpr bool kIsVariableC = true; + BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] { + BOOL_SWITCH(params.z_ptr != nullptr , kHasZ, [&] { + BOOL_SWITCH(params.query_start_loc_ptr != nullptr , kVarlen, [&] { + using Ktraits = Selective_Scan_fwd_kernel_traits; + constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * MAX_DSTATE * sizeof(typename Ktraits::scan_t); + dim3 grid(params.batch, params.dim / kNRows); + auto kernel = &selective_scan_fwd_kernel; + if (kSmemSize >= 48 * 1024) { +#ifdef USE_ROCM + STD_CUDA_CHECK(hipFuncSetAttribute( + reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); +#else + STD_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); +#endif + } + kernel<<>>(params); + STD_CUDA_KERNEL_LAUNCH_CHECK(); + }); + }); + }); +} + +template +void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) { + + #ifndef USE_ROCM + if (params.cache_enabled && params.block_size == 1024) { + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 128) { + selective_scan_fwd_launch<32, 4, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 256) { + selective_scan_fwd_launch<32, 8, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 512) { + selective_scan_fwd_launch<32, 16, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 1024) { + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); + } else { + selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); + } + #else + if (params.cache_enabled && params.block_size == 1024) { + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 256) { + selective_scan_fwd_launch<64, 4, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 512) { + selective_scan_fwd_launch<64, 8, input_t, weight_t, state_t>(params, stream); + } else if (params.seqlen <= 1024) { + selective_scan_fwd_launch<64, 16, input_t, weight_t, state_t>(params, stream); + } else { + selective_scan_fwd_launch<128, 16, input_t, weight_t, state_t>(params, stream); + } + #endif +} diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_bf16.cu b/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_bf16.cu new file mode 100644 index 000000000000..d2fcf9abb4f3 --- /dev/null +++ b/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_bf16.cu @@ -0,0 +1,10 @@ +// Explicit instantiations for bf16 input type. + +#include "kernel.cuh" + +template void selective_scan_fwd_cuda( + SSMParamsBase& params, cudaStream_t stream); +template void selective_scan_fwd_cuda(SSMParamsBase& params, + cudaStream_t stream); diff --git a/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_fp16.cu b/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_fp16.cu new file mode 100644 index 000000000000..53cdca294533 --- /dev/null +++ b/csrc/libtorch_stable/mamba/selective_scan_fwd/shard_fp16.cu @@ -0,0 +1,11 @@ +// Explicit instantiations for fp16 and fp32 input types. + +#include "kernel.cuh" + +template void selective_scan_fwd_cuda( + SSMParamsBase& params, cudaStream_t stream); +template void selective_scan_fwd_cuda( + SSMParamsBase& params, cudaStream_t stream); +template void selective_scan_fwd_cuda( + SSMParamsBase& params, cudaStream_t stream); diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 9e6c0726d19e..36bcd0cc9ad3 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -16,3608 +16,31 @@ #include #include -#include -#include -#include -#include "../cuda_compat.h" -#include -#include "../attention/dtype_fp8.cuh" -#include "../quantization/w8a8/fp8/amd/quant_utils.cuh" +#include -// ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent -#if !defined(HIP_FP8_TYPE_OCP) -using __hip_fp8_e4m3 = __hip_fp8_e4m3_fnuz; -using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz; -#endif - -#if defined(__HIPCC__) && \ - (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) - #define __HIP__GFX9__ -#endif - -#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) - #define __HIP__FP8MFMA__ -#endif - -#if defined(NDEBUG) - #undef NDEBUG - #include - #define UNREACHABLE_CODE assert(false); - #define NDEBUG -#else - #define UNREACHABLE_CODE assert(false); -#endif - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -enum class MFMAType { - F16 = 0, - Fp8 = 1, - Fp4 = 2, -}; - -#if defined(__HIP__GFX9__) - - #define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32 - #define GCN_MFMA_INSTR __builtin_amdgcn_mfma_f32_4x4x4f16 - -using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float; -using float16x4 = - __attribute__((__vector_size__(4 * sizeof(_Float16)))) _Float16; -typedef float16x4 _Half4; -using float16x2 = - __attribute__((__vector_size__(2 * sizeof(_Float16)))) _Float16; -typedef float16x2 _Half2; -typedef struct _Half8 { - _Half4 xy[2]; -} _Half8; - -using bit16_t = uint16_t; -using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; -typedef bit16x4 _B16x4; -typedef struct _B16x8 { - _B16x4 xy[2]; -} _B16x8; - -using _B8x8 = uint2; -using _B8x4 = int32_t; // used in builtins -using bit8_t = uint8_t; - -typedef struct _B8x16 { - _B8x8 xy[2]; -} _B8x16; - -template -__device__ __forceinline__ floatx4 gcn_mfma4x4x4_instr(const _B16x4& inpA, - const _B16x4& inpB, - const floatx4& inpC) { - if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_4x4x4f16(inpA, inpB, inpC, absz, cbid, - blgp); - } else if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(inpA, inpB, inpC, absz, cbid, - blgp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA, - const _B16x4& inpB, - const floatx4& inpC) { - if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_16x16x16f16(inpA, inpB, inpC, absz, cbid, - blgp); - } else if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(inpA, inpB, inpC, absz, - cbid, blgp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ floatx4 gcn_mfma16x16x32_instr(const long& inpA, - const long& inpB, - const floatx4& inpC) { - if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(inpA, inpB, inpC, absz, - cbid, blgp); - } else if constexpr (std::is_same::value) { - return __builtin_amdgcn_mfma_f32_16x16x32_bf8_bf8(inpA, inpB, inpC, absz, - cbid, blgp); - } else { - static_assert(false, "unsupported 8b dtype"); - } -} - -template -__device__ __forceinline__ float to_float(const T& inp) { - if constexpr (std::is_same::value) { - return (float)inp; - } else if constexpr (std::is_same::value) { - return __bfloat162float(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ T from_float(const float& inp) { - if constexpr (std::is_same::value) { - return (_Float16)inp; - } else if constexpr (std::is_same::value) { - return __float2bfloat16(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ _B16x4 from_floatx4(const floatx4& inp) { - _B16x4 ret; - if constexpr (std::is_same::value) { - union h2cvt { - __half2 h2[2]; - _B16x4 b16x4; - } u; - u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); - u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); - return u.b16x4; - } else if constexpr (std::is_same::value) { - for (int i = 0; i < 4; i++) { - union fcvt { - uint32_t u32; - float f32; - } u; - u.f32 = inp[i]; - u.u32 += 0x7fff + ((u.u32 >> 16) & 1); // BF16 RNE with no nan/inf check - ret[i] = uint16_t(u.u32 >> 16); - } - return ret; - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ _B16x4 addx4(const _B16x4& inp1, - const _B16x4& inp2) { - _B16x4 ret; - if constexpr (std::is_same::value) { - union h2cvt { - _B16x4 b16x4; - __half2 h2[2]; - } u1, u2, s; - u1.b16x4 = inp1; - u2.b16x4 = inp2; - s.h2[0] = u1.h2[0] + u2.h2[0]; - s.h2[1] = u1.h2[1] + u2.h2[1]; - return s.b16x4; - } else if constexpr (std::is_same::value) { - for (int i = 0; i < 4; i++) { - union fcvt { - float f32; - uint32_t i32; - } u1, u2, s; - u1.i32 = uint32_t(inp1[i]) << 16; - u2.i32 = uint32_t(inp2[i]) << 16; - s.f32 = u1.f32 + u2.f32; - ret[i] = uint16_t(s.i32 >> 16); - } - return ret; - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -__device__ __forceinline__ floatx4 to_float_fp8x4(const _B8x4& inp) { - // From MI300+ platforms, we have v_cvt_pk_f32_fp8 instruction - // to convert 2 packed fp8 to 2 packed fp32 values. - // However, in MI200 platforms, we only have v_cvt_f32_fp8 - // to convert fp8 values individually. So we added - // #else case for fewer instructions (# inst=2) in MI300+, - // and fallback to - // #if case for other platforms (# inst=4). - #if defined(__gfx90a__) - float4 f32x4 = vllm::fp8::vec_conversion( - *reinterpret_cast(&inp)); - return *reinterpret_cast(&f32x4); - #else // MI3xx+ optimized builtins - const auto f0 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, false); - const auto f1 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, true); - floatx4 ret; - ret[0] = f0[0]; - ret[1] = f0[1]; - ret[2] = f1[0]; - ret[3] = f1[1]; - return ret; - #endif -} - -template -__device__ __forceinline__ _B16x4 from_floatx4_rtz(const floatx4& inp) { - _B16x4 ret; - if constexpr (std::is_same::value) { - union h2cvt { - _Half2 h2[2]; - _B16x4 b16x4; - } u; - u.h2[0] = __builtin_amdgcn_cvt_pkrtz(inp[0], inp[1]); - u.h2[1] = __builtin_amdgcn_cvt_pkrtz(inp[2], inp[3]); - return u.b16x4; - } else if constexpr (std::is_same::value) { - for (int i = 0; i < 4; i++) { - union fcvt { - uint32_t i32; - float f32; - } u; - u.f32 = inp[i]; - ret[i] = uint16_t(u.i32 >> 16); - } - return ret; - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) { - union { - _B8x8 b8x8; - _B8x4 b8x4[2]; - } tmp; - tmp.b8x8 = input; - _B16x8 ret; - for (int i = 0; i < 2; i++) { - ret.xy[i] = from_floatx4_rtz(to_float_fp8x4(tmp.b8x4[i])); - } - return ret; -} - -typedef union u64_cvt { - half f16x4[4]; - int16_t b16x4[4]; - _B8x8 b8x8; - _B16x4 b64; - int64_t i64; -} _T8x8; - -__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input, - _T8x8& Mtemp) { - _T8x8 Qtmp8x8; - - for (int i = 0; i < 2; i++) { - floatx4 q_out = {0, 0, 0, 0}; - q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i], - q_out); - Qtmp8x8.b16x4[i * 2] = - __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false); - Qtmp8x8.b16x4[i * 2 + 1] = - __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false); - } - return Qtmp8x8.b8x8; -} - -__device__ float warpReduceMax(float val) { - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - val = max( - val, __shfl_down(val, offset, WARP_SIZE)); // Using max() for reduction - } - return val; -} - -// grid (num_seqs, num_partitions,num_kv_heads) -// block (256) -// clang-format off -template -__global__ -__launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] - const int num_kv_heads, - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, - const int kv_block_stride, - const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - // clang-format on - constexpr int NWARPS = NUM_THREADS / WARP_SIZE; - const auto warpid = threadIdx.x / WARP_SIZE; - const auto laneid = threadIdx.x % WARP_SIZE; - const int lane4id = laneid % 4; - const int lane16id = laneid % 16; - const int rowid = laneid / 16; - - const auto seq_idx = blockIdx.x; - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { - return; - } - - const auto partition_idx = blockIdx.y; - - constexpr int T_PAR_SIZE = 256; // token partition size set to 256 - - const auto max_num_partitions = gridDim.y; - - const int seq_len = seq_lens[seq_idx]; - - const int partition_start_token_idx = - partition_idx * T_PAR_SIZE; // partition_size; - // exit if partition is out of context for seq - if (partition_start_token_idx >= seq_len) { - return; - } - - constexpr int GQA_RATIO4 = DIVIDE_ROUND_UP(GQA_RATIO, 4); - - // shared_logits is used for multiple purposes - __shared__ _B16x4 shared_logits[NWARPS][4][16][4]; - - // for QK mfma16x16, layout is QHead/Tokenx16 across every 16 lanes, 16 Bytes - // HeadElements in each lane, 4x16B HeadElements across 4 rows of warp - constexpr int ROWS_PER_WARP = - WARP_SIZE / 16; // rows refers to 16 lanes; refer DDP (Data Parallel - // Processing) terminology - constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = - 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types - constexpr int QKHE_PER_FETCH = - CONTIGUOUS_KV_ELEMS_16B_LOAD * - ROWS_PER_WARP; // each fetch across a warp fetches these many elements - constexpr int QK_SIZE_RATIO = - sizeof(scalar_t) / - sizeof(cache_t); // 1 for 16bit types, 2 for 8bit types - constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 4xQKHE_16B across - // warp - - _B16x8 Qlocal[QKHELOOP] - [QK_SIZE_RATIO]; // note that 16 contiguous elements of Q should - // be fetched per lane for 8 bit cache types : - // QK_SIZE_RATIO changes for this - - constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); - - constexpr int TOKENS_PER_WARP = - T_PAR_SIZE / - NWARPS; // sub partition of tokens per warp for qk calculation - constexpr int TLOOP = - TOKENS_PER_WARP / - 16; // each mfma16x16x16 instruction processes 16 tokens - - // can be interpreted as B8x16 for 8 bit types - _B16x8 Klocal[TLOOP][QKHELOOP]; - - const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; - const auto wg_start_kv_head_idx = blockIdx.z; - const auto total_num_heads = gridDim.z * GQA_RATIO; - - // for QK mfma, tokens in multiples of TOKENS_PER_WARP are spread across warps - // each mfma takes QH16xT16x16HE across warp - // repeat mfmas across QKHELOOP dimension - // output layout from QKmfma : QH16xT4x4 16 qheads across 16 lanes, 16 tokens - // across 4 rows x 4 tokens per lane - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int last_seq_block = num_seq_blocks - 1; - - const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; - - int kphysical_block_number[TLOOP]; - #if defined(__HIP__FP8MFMA__) - float q_max = 0; - float q_scale = 1.0; - #endif - - // fetch k physical block numbers - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; - const int kblock_idx = (kglobal_token_idx < seq_len) - ? kglobal_token_idx / BLOCK_SIZE - : last_seq_block; - kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; - } - - // fetch Q in shared across warps and then write to registers - const int local_qhead_idx = 4 * warpid + rowid; - const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - const scalar_t* q_ptr = - q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; - - const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; - if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { - const scalar_t* q_fetch_ptr = q_ptr + qhead_element; - const _B16x8* q_fetch_ptr_16B = - reinterpret_cast(q_fetch_ptr); - _B16x8 tmp = *q_fetch_ptr_16B; - if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { - const int offset1 = - lane16id / - 4; // 16 contiguous chunks of head elems are spread across 4x4lanes - shared_logits[offset1][lane4id][local_qhead_idx][0] = tmp.xy[0]; - shared_logits[offset1][lane4id][local_qhead_idx][1] = tmp.xy[1]; - } else { - for (int i = 0; i < 2; i++) { - const int head_elem = lane16id * 2 + i; // element id in _B16x4 terms - const int offset3 = head_elem % 4; - const int offset2 = (head_elem / 4) % 4; - const int offset1 = head_elem / 4 / 4; - shared_logits[offset1][offset2][local_qhead_idx][offset3] = tmp.xy[i]; - } - } - } - __syncthreads(); - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { - for (int i = 0; i < 2; i++) { - Qlocal[qkhe_depth][qkratio].xy[i] = - shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO] - [2 * qkratio + i]; - #if defined(__HIP__FP8MFMA__) - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto && - MFMA_TYPE == MFMAType::Fp8) { - scalar_t* qptr = - reinterpret_cast(&Qlocal[qkhe_depth][qkratio].xy[i]); - for (int k = 0; k < 4; k++) - q_max = fmax(fabs(to_float(qptr[k])), q_max); - } - #endif - } - } - } - - constexpr int KX = - 16 / sizeof(cache_t); // vLLM defines x as 16 Bytes of kv cache elements - const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; - - const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; - // fetch K values - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int64_t kblock_number = - static_cast(kphysical_block_number[token_depth]); - const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; - const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; - - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; - const int offset1 = head_elem / KX; - const int offset2 = head_elem % KX; - const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; - const _B16x8* k_fetch_ptr_16B = - reinterpret_cast(k_fetch_ptr); - Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; - } - } - - float alibi_slope; - if constexpr (ALIBI_ENABLED) { - const int alibi_head_idx = wg_start_head_idx + lane16id; - alibi_slope = (lane16id < GQA_RATIO) ? alibi_slopes[alibi_head_idx] : 0.f; - } - - constexpr int VTOKENS_PER_LANE = - TOKENS_PER_WARP / ROWS_PER_WARP; // 64/4 = 16 contiguous vtokens per lane - constexpr int VBLOCKS_PER_LANE = - 1; // assumes block size >=16, each lane can correspond to 1 block only - constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps - constexpr int VTLANELOOP = DIVIDE_ROUND_UP( - VTOKENS_PER_LANE, - CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes - // minimum block size is 16 - constexpr int VHELOOP = HEAD_SIZE / 16 / NWARPS; - - int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; - - // fetch v physical block numbers - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; - vblock_depth++) { - const int vlocal_token_idx = - vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + - rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; - // Safe to use an int32_t here assuming we are working with < 2 billion - // tokens - const int vglobal_token_idx = - partition_start_token_idx + vlocal_token_idx; - const int vblock_idx = (vglobal_token_idx < seq_len) - ? vglobal_token_idx / BLOCK_SIZE - : last_seq_block; - vphysical_block_number[vtoken_depth][vblock_depth] = - block_table_seq[vblock_idx]; - } - } - - _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too - - const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + - ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); - - // v fetches are 16head elems across lanes x 16 tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; - const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - const int vblock_depth = 0; - const int64_t vblock_number = static_cast( - vphysical_block_number[vtoken_depth][vblock_depth]); - const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); - - const cache_t* v_fetch_ptr = - v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; - const _B16x8* v_fetch_ptr_16B = - reinterpret_cast(v_fetch_ptr); - Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; - } - } - } - - // calculate post qk mfma scale - float scale2 = scale; - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { - // multiply by k_scale if fp8 kv cache - scale2 *= *k_scale; - #if defined(__HIP__FP8MFMA__) - q_max = warpReduceMax(q_max); - constexpr float FP8_E4M3_SCALE_TARGET = 224.0f; - if constexpr (MFMA_TYPE == MFMAType::Fp8) { - q_scale = q_max > 0 ? FP8_E4M3_SCALE_TARGET / q_max : 1.0f; - scale2 /= q_scale; - } - #endif - } - - floatx4 d_out[TLOOP]; - // qk mfma - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - d_out[token_depth] = {0}; - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { - for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { - for (int i = 0; i < 2; i++) { - d_out[token_depth] = gcn_mfma16x16x16_instr( - Klocal[token_depth][qkhe_depth].xy[i], - Qlocal[qkhe_depth][qkratio].xy[i], d_out[token_depth]); - } - } - } else { // kv cache dtype fp8 - auto Ktmp = Klocal[token_depth][qkhe_depth]; - _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp); - for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { - if constexpr (MFMA_TYPE == MFMAType::F16) { - _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio]; - _B16x8 Klocaltmp = convert_b8x8_custom(Ktmp8x8); - for (int i = 0; i < 2; i++) { - d_out[token_depth] = gcn_mfma16x16x16_instr( - Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i], - d_out[token_depth]); - } - } else { - #if defined(__HIP__FP8MFMA__) - _T8x8 Ktmp8x8, Qtmp8x8; - Ktmp8x8.b8x8 = Ktmp8x16.xy[qkratio]; - - for (int n = 0; n < 2; n++) { - scalar_t* qptr = reinterpret_cast( - &Qlocal[qkhe_depth][qkratio].xy[n]); - - Qtmp8x8.b16x4[n * 2] = - vllm::fp8::scaled_vec_conversion( - make_float2(to_float(qptr[0]), - to_float(qptr[1])), - q_scale); - Qtmp8x8.b16x4[n * 2 + 1] = - vllm::fp8::scaled_vec_conversion( - make_float2(to_float(qptr[2]), - to_float(qptr[3])), - q_scale); - } - - d_out[token_depth] = - gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( - Ktmp8x8.i64, Qtmp8x8.i64, d_out[token_depth]); - #else - UNREACHABLE_CODE - #endif - } - } - } - } - d_out[token_depth] *= scale2; - } - - const int qkout_token_idx = - partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 4; - - // apply alibi - if constexpr (ALIBI_ENABLED) { - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - const int alibi_offset = local_token_idx - seq_len + 1; - for (int i = 0; i < 4; i++) { - d_out[token_depth][i] += alibi_slope * (alibi_offset + i); - } - } - } - - // calculate qk_max and exp_sum per warp and write to shared memory - float qk_max = -FLT_MAX; - float exp_sum = 0.0f; - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 4; i++) { - const float tmp = - (local_token_idx + i < seq_len) ? d_out[token_depth][i] : -FLT_MAX; - qk_max = fmaxf(qk_max, tmp); - } - } - - for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { - qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask)); - } - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 4; i++) { - const float tmp = (local_token_idx + i < seq_len) - ? __expf(d_out[token_depth][i] - qk_max) - : 0.0f; - d_out[token_depth][i] = tmp; - exp_sum += tmp; - } - } - - for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { - exp_sum += __shfl_xor(exp_sum, mask); - } - - __syncthreads(); // sync before writing to shared mem - - float* shared_mem = reinterpret_cast(shared_logits); - if (laneid < 16) { - const int qk_max_offset = warpid * 16 + lane16id; - shared_mem[qk_max_offset] = qk_max; - const int exp_sum_offset = NWARPS * 16 + qk_max_offset; - shared_mem[exp_sum_offset] = exp_sum; - } - - __syncthreads(); - - // calculate partition qk_max and exp_sum - float partition_qk_max = -FLT_MAX; - float warp_qk_max_exp[NWARPS]; - float partition_exp_sum = 0.0f; - - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = shared_mem[w * 16 + lane16id]; - partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); - } - - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); - partition_exp_sum += - shared_mem[NWARPS * 16 + w * 16 + lane16id] * warp_qk_max_exp[w]; - } - - const float inv_sum_scale = - __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; - - __syncthreads(); - - // disable rtz conversion due to its impact on accuracy. - constexpr bool LOGITS_RTZ_CONVERSION = false; - - #if defined(__HIP__FP8MFMA__) - int rowid_8x8 = rowid / 2; - int offset = rowid % 2; - #endif - - // write logits to shared mem - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - d_out[token_depth] *= inv_sum_scale; - if constexpr (MFMA_TYPE != MFMAType::Fp8) { - if constexpr (LOGITS_RTZ_CONVERSION) { - // use rtz conversion for better performance, with negligible impact on - // accuracy - shared_logits[warpid][token_depth][lane16id][rowid] = - from_floatx4_rtz(d_out[token_depth]); - } else { - shared_logits[warpid][token_depth][lane16id][rowid] = - from_floatx4(d_out[token_depth]); - } - } else { - #if defined(__HIP__FP8MFMA__) - // cast _B16x4* to _B8x8* - _T8x8& logits_8x8 = *reinterpret_cast<_T8x8*>( - &shared_logits[warpid][token_depth][lane16id][rowid_8x8]); - logits_8x8.b16x4[offset * 2] = __builtin_amdgcn_cvt_pk_fp8_f32( - d_out[token_depth][0], d_out[token_depth][1], 0, false); - logits_8x8.b16x4[offset * 2 + 1] = __builtin_amdgcn_cvt_pk_fp8_f32( - d_out[token_depth][2], d_out[token_depth][3], 0, false); - #else - UNREACHABLE_CODE - #endif - } - } - - // write out partition max_logits and exp_sum - if (threadIdx.x < GQA_RATIO) { - const int qhead_idx = lane16id; - const int64_t offset = static_cast(seq_idx) * - static_cast(total_num_heads) * - static_cast(max_num_partitions) + - (static_cast(wg_start_head_idx) + - static_cast(qhead_idx)) * - static_cast(max_num_partitions) + - static_cast(partition_idx); - max_logits[offset] = partition_qk_max; - exp_sums[offset] = partition_exp_sum; - } - - __syncthreads(); - - constexpr int ELEMS8_ELEMS4_RATIO = 8 / 4; - constexpr int ELEMS16_ELEMS8_RATIO = 16 / 8; - - _B16x4 outelems[VHELOOP]; - // Softmax V mfma - // v layout: 16he across lanes x 16 tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - floatx4 tmp_out = {0}; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { - const int offset = rowid * VTLANELOOP * ELEMS8_ELEMS4_RATIO + - vfetch_depth * ELEMS8_ELEMS4_RATIO + i; - const int offset1 = offset % ROWS_PER_WARP; - const int offset2 = offset / ROWS_PER_WARP; - // output format is 16 qheads across 16 lanes, 16 head elems spread - // across 4 rows - tmp_out = gcn_mfma16x16x16_instr( - Vlocal[vtoken_depth][vhe_depth][vfetch_depth].xy[i], - shared_logits[vtoken_depth][offset2][lane16id][offset1], - tmp_out); - } - } - // KV cache fp8 - } else { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - _B16x8 Vtmp = Vlocal[vtoken_depth][vhe_depth][vfetch_depth]; - // reinterpret V format as 16 elements of 8bits - _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp); - for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) { - _B8x8 Vtmp8x8 = Vtmp8x16.xy[j]; - if constexpr (MFMA_TYPE == MFMAType::F16) { - _B16x8 Vlocaltmp = convert_b8x8_custom(Vtmp8x8); - for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { - const int offset = - rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + - j * ELEMS8_ELEMS4_RATIO + i; - const int offset1 = offset % ROWS_PER_WARP; - const int offset2 = offset / ROWS_PER_WARP; - // output format is 16 qheads across 16 lanes, 16 head elems - // spread across 4 rows - tmp_out = gcn_mfma16x16x16_instr( - Vlocaltmp.xy[i], - shared_logits[vtoken_depth][offset2][lane16id][offset1], - tmp_out); - } - } else { - #if defined(__HIP__FP8MFMA__) - for (int i = 0; i < ELEMS8_ELEMS4_RATIO / 2; i++) { - const int offset = - rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + - j * ELEMS8_ELEMS4_RATIO + i; - const int offset1 = (offset % ROWS_PER_WARP) / 2; - const int offset2 = offset / ROWS_PER_WARP; - // output format is 16 qheads across 16 lanes, 16 head elems - // spread across 4 rows - tmp_out = gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( - reinterpret_cast<_T8x8*>(&Vtmp8x8)->i64, - reinterpret_cast<_T8x8*>( - &shared_logits[vtoken_depth][offset2][lane16id] - [offset1]) - ->i64, - tmp_out); - } - #else - UNREACHABLE_CODE - #endif - } - } - } - } - } - // apply post Softmax V mfma v_scale - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { - tmp_out *= *v_scale; - } - outelems[vhe_depth] = from_floatx4(tmp_out); - } - - __syncthreads(); - - // store Softmax-V mfma output to shared mem - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - // lane16 id head dimension; rowid head element dimension - shared_logits[warpid][vhe_depth][lane16id][rowid] = outelems[vhe_depth]; - } - - __syncthreads(); - - // write to tmp_out with coalesced writes after reading from shared mem - if (warpid == 0) { - _B16x8 vout[GQA_RATIO4]; - // each lane writes out 16Bytes of tmp_out along head elem dimension - const int head_elem_idx = lane16id * 8; - if (head_elem_idx < HEAD_SIZE) { - for (int h = 0; h < GQA_RATIO4; h++) { - const int local_head_idx = 4 * h + rowid; - const int offset1 = (head_elem_idx / 16) % 4; - const int offset2 = head_elem_idx / 16 / NWARPS; - const int offset3 = (head_elem_idx / 4) % 4; - for (int i = 0; i < 2; i++) { - vout[h].xy[i] = - shared_logits[offset1][offset2][local_head_idx][offset3 + i]; - } - } - - const int64_t hsz_maxp_mult = - static_cast(HEAD_SIZE * max_num_partitions); - scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + - partition_idx * HEAD_SIZE; - for (int h = 0; h < GQA_RATIO4; h++) { - const int local_head_idx = 4 * h + rowid; - if (local_head_idx < GQA_RATIO) { - const int64_t out_head_idx = - static_cast(wg_start_head_idx + local_head_idx); - scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; - scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; - _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); - *out_ptr_B16x8 = vout[h]; - } - } - } - } -} - -// grid (num_seqs, num_partitions, num_kv_heads) -// block (256 : partition size) -// each WG handles 1 partition per sequence -// clang-format off -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] - const int num_kv_heads, - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, - const int kv_block_stride, - const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - // clang-format on - constexpr int NWARPS = NUM_THREADS / WARP_SIZE; - const auto warpid = threadIdx.x / WARP_SIZE; - const auto laneid = threadIdx.x % WARP_SIZE; - const int lane4id = laneid % 4; - - const auto seq_idx = blockIdx.x; - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { - return; - } - const auto partition_idx = blockIdx.y; - const auto partition_size = blockDim.x; - const auto max_num_partitions = gridDim.y; - - const int seq_len = seq_lens[seq_idx]; - const int partition_start_token_idx = partition_idx * partition_size; - // exit if partition is out of context for seq - if (partition_start_token_idx >= seq_len) { - return; - } - // every 4 lanes fetch 4 different qheads - // qhloop = num loops over qhead dimension - constexpr int QHLOOP = DIVIDE_ROUND_UP(GQA_RATIO, 4); - constexpr int GQA_RATIO4 = 4 * QHLOOP; - __shared__ float shared_qk_max[NWARPS][GQA_RATIO4 + 1]; - __shared__ float shared_exp_sum[NWARPS][GQA_RATIO4 + 1]; - _B16x8 Qlocal[QHLOOP]; - constexpr int x = 16 / sizeof(scalar_t); - // kheloop = num loops over head_size for 16Bytes of Q/dequantized K elements - constexpr int KHELOOP = HEAD_SIZE / x; - _B16x8 Klocal[KHELOOP]; - _B8x8 Klocalb8[KHELOOP]; - // for SoftMax-V Gemm, V head_size dimension is distributed across warp - // vheloop = num loops to cover v head size dimension - constexpr int VHELOOP = HEAD_SIZE / WARP_SIZE; - // softmax out has warp_size tokens across warp - // vtloop = num loops to cover warp_size(64) tokens with 16Bytes of - // dequantized V elements - constexpr int VTLOOP = WARP_SIZE / 8; - // num vblocks to cover warp_size(64) v elements - constexpr int VBLOCKS = 8 * VTLOOP / BLOCK_SIZE; - int vphysical_blocks[VBLOCKS]; - _B16x8 Vlocal[VHELOOP][VTLOOP]; - _B8x8 Vlocalb8[VHELOOP][VTLOOP]; - floatx4 d_out[QHLOOP]; - float qk_max[QHLOOP]; - - __shared__ _B16x4 vout_shared[QHLOOP][VHELOOP][WARP_SIZE][NWARPS + 1]; - - for (int h = 0; h < QHLOOP; h++) { - d_out[h] = {0}; - qk_max[h] = -FLT_MAX; - } - - const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; - const auto wg_start_kv_head_idx = blockIdx.z; - - const int warp_start_token_idx = - partition_start_token_idx + warpid * WARP_SIZE; - - if (warp_start_token_idx >= seq_len) { // warp out of context - #pragma unroll - for (int h = 0; h < GQA_RATIO4; h++) { - shared_qk_max[warpid][h] = -FLT_MAX; - shared_exp_sum[warpid][h] = 0.0f; - } - } else { // warp within context - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int last_seq_block = num_seq_blocks - 1; - - const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; - // token id within partition - const auto local_token_idx = threadIdx.x; - // token id within sequence - const int global_token_idx = partition_start_token_idx + local_token_idx; - - // fetch block number for k - const int block_idx = (global_token_idx < seq_len) - ? global_token_idx / BLOCK_SIZE - : last_seq_block; - - // fetch k physical block number - // int32 physical_block_number leads to overflow when multiplied with - // kv_block_stride - const int64_t physical_block_number = - static_cast(block_table[block_idx]); - - // fetch vphysical block numbers up front - const int warp_start_block_idx = warp_start_token_idx / BLOCK_SIZE; - for (int b = 0; b < VBLOCKS; b++) { - const int vblock_idx = warp_start_block_idx + b; - const int vblock_idx_ctx = - (vblock_idx <= last_seq_block) ? vblock_idx : last_seq_block; - vphysical_blocks[b] = block_table[vblock_idx_ctx]; - } - - // fetch q elements - // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elemsc - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - const scalar_t* q_ptr = - q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE; - const _B16x8* q_ptrh8 = reinterpret_cast(q_ptr); - const int qhead_elemh8 = laneid / 4; - - for (int h = 0; h < QHLOOP - 1; h++) { - const int qhead_idx = h * 4 + lane4id; - Qlocal[h] = q_ptrh8[qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; - } - const int final_qhead_idx = 4 * (QHLOOP - 1) + lane4id; - if (final_qhead_idx < GQA_RATIO) { - Qlocal[QHLOOP - 1] = - q_ptrh8[final_qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; - } else { - Qlocal[QHLOOP - 1].xy[0] = {0}; - Qlocal[QHLOOP - 1].xy[1] = {0}; - } - - // fetch k elements - const cache_t* k_ptr = k_cache + physical_block_number * kv_block_stride + - wg_start_kv_head_idx * kv_head_stride; - - // physical_block_offset is already cast in terms of _B16x8 - const int physical_block_offset = local_token_idx % BLOCK_SIZE; - - // each K fetch is for 8 elements of cache_t which are later dequantized to - // scalar_t for fp8 - if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { - const _B16x8* k_ptrh8 = reinterpret_cast(k_ptr); - for (int d = 0; d < KHELOOP; d++) { - Klocal[d] = k_ptrh8[d * BLOCK_SIZE + physical_block_offset]; - } - } else { - // vllm defines X as 16 Bytes of elements of cache_t - constexpr int X = 16 / sizeof(cache_t); - const cache_t* k_ptr2 = k_ptr + physical_block_offset * X; - for (int d = 0; d < KHELOOP; d++) { - const int head_elem = d * 8; - const int offset1 = head_elem / X; - const int offset2 = head_elem % X; - const cache_t* k_ptr3 = k_ptr2 + offset1 * BLOCK_SIZE * X + offset2; - Klocalb8[d] = *reinterpret_cast(k_ptr3); - } - } - - // optional alibi fetch - float alibi_slope[QHLOOP]; - if constexpr (ALIBI_ENABLED) { - for (int h = 0; h < QHLOOP; h++) { - const int qhead_idx = h * 4 + lane4id; - alibi_slope[h] = (qhead_idx < GQA_RATIO) - ? alibi_slopes[wg_start_head_idx + qhead_idx] - : 0.f; - } - } - - const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; - // fetch vcache in kv cache auto case - if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { - const _B16x8* v_ptrh8 = reinterpret_cast(v_ptr); - // iterate over each v block - for (int b = 0; b < VBLOCKS; b++) { - // int32 physical_block_number leads to overflow when multiplied with - // kv_block_stride - const int64_t vphysical_block_number = - static_cast(vphysical_blocks[b]); - const _B16x8* v_ptrh8b = - v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; - // iterate over each head elem (within head_size) - for (int h = 0; h < VHELOOP; h++) { - const int head_size_elem = h * WARP_SIZE + laneid; - const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; - // iterate over all velems within block - for (int d = 0; d < BLOCK_SIZE / 8; d++) { - Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; - } - } - } - } // if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) - // fetch vcache in fp8 case - else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) - const _B8x8* v_ptrh8 = reinterpret_cast(v_ptr); - // iterate over each v block - for (int b = 0; b < VBLOCKS; b++) { - // int32 physical_block_number leads to overflow when multiplied with - // kv_block_stride - const int64_t vphysical_block_number = - static_cast(vphysical_blocks[b]); - const _B8x8* v_ptrh8b = - v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; - // iterate over each head elem (within head_size) - for (int h = 0; h < VHELOOP; h++) { - const int head_size_elem = h * WARP_SIZE + laneid; - const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; - // iterate over all velems within block - for (int d = 0; d < BLOCK_SIZE / 8; d++) { - Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; - } - } - } - } - - #define QK_mfma(x) \ - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ - Klocal[x] = convert_b8x8_custom(Klocalb8[x]); \ - } \ - for (int h = 0; h < QHLOOP; h++) { \ - d_out[h] = gcn_mfma4x4x4_instr( \ - Qlocal[h].xy[0], Klocal[x].xy[0], d_out[h]); \ - d_out[h] = gcn_mfma4x4x4_instr( \ - Qlocal[h].xy[1], Klocal[x].xy[1], d_out[h]); \ - } - // QK mfma with Q mfma block broadcast - // Q values across head_size dimension stored across lanes - // K values across head_size dimension are stored depthwise within lane - // Q broadcast with absz, cbid of mfma instruction - QK_mfma(0); - QK_mfma(1); - QK_mfma(2); - QK_mfma(3); - QK_mfma(4); - QK_mfma(5); - QK_mfma(6); - QK_mfma(7); - // below only needed for head size 128 - if constexpr (KHELOOP > 8) { - QK_mfma(8); - QK_mfma(9); - QK_mfma(10); - QK_mfma(11); - QK_mfma(12); - QK_mfma(13); - QK_mfma(14); - QK_mfma(15); - } - #undef QK_mfma - - float scale2 = scale; - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { - // post mfma scaling for fp8 - scale2 *= *k_scale; - } - - for (int h = 0; h < QHLOOP; h++) { - d_out[h] *= scale2; - } - - // transpose d_out so that 4 token ids are in each lane, and 4 heads are - // across 4 lanes - for (int h = 0; h < QHLOOP; h++) { - floatx4 tmp = {0}; - for (int i = 0; i < 4; i++) { - const float B = (lane4id == i) ? 1.0f : 0.0f; - tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(d_out[h][i], B, tmp, 0, 0, 0); - } - d_out[h] = tmp; - } - - const int lane4_token_idx = 4 * (global_token_idx >> 2); - - if constexpr (ALIBI_ENABLED) { - const int alibi_offset = lane4_token_idx - seq_len + 1; - for (int h = 0; h < QHLOOP; h++) { - for (int i = 0; i < 4; i++) { - d_out[h][i] += alibi_slope[h] * (alibi_offset + i); - } - } - } - - const int bpermute_mask = 4 * (16 * ((laneid >> 2) % 4) + lane4id); - - for (int h = 0; h < QHLOOP; h++) { - qk_max[h] = -FLT_MAX; - for (int i = 0; i < 4; i++) { - qk_max[h] = (lane4_token_idx + i < seq_len) - ? fmaxf(qk_max[h], d_out[h][i]) - : qk_max[h]; - } - - // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { - // qk_max[h] = fmaxf(qk_max[h], __shfl_xor(qk_max[h], mask)); - // } - // faster version of above code with dpp - asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" - : "=v"(qk_max[h]) - : "v"(qk_max[h]), "v"(qk_max[h])); - asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" - : "=v"(qk_max[h]) - : "v"(qk_max[h]), "v"(qk_max[h])); - - auto tmp = __builtin_amdgcn_ds_bpermute( - bpermute_mask, *reinterpret_cast(&qk_max[h])); - qk_max[h] = *reinterpret_cast(&tmp); - asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" - : "=v"(qk_max[h]) - : "v"(qk_max[h]), "v"(qk_max[h])); - asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" - : "=v"(qk_max[h]) - : "v"(qk_max[h]), "v"(qk_max[h])); - } - - float exp_sum[QHLOOP]; - for (int h = 0; h < QHLOOP; h++) { - exp_sum[h] = 0.0f; - for (int i = 0; i < 4; i++) { - d_out[h][i] = (lane4_token_idx + i < seq_len) - ? __expf(d_out[h][i] - qk_max[h]) - : 0.0f; - exp_sum[h] += d_out[h][i]; - } - // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { - // exp_sum[h] += __shfl_xor(exp_sum[h], mask); - // } - // faster version of above code with dpp - asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" - : "=v"(exp_sum[h]) - : "v"(exp_sum[h]), "v"(exp_sum[h])); - asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" - : "=v"(exp_sum[h]) - : "v"(exp_sum[h]), "v"(exp_sum[h])); - - auto tmp = __builtin_amdgcn_ds_bpermute( - bpermute_mask, *reinterpret_cast(&exp_sum[h])); - exp_sum[h] = *reinterpret_cast(&tmp); - asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" - : "=v"(exp_sum[h]) - : "v"(exp_sum[h]), "v"(exp_sum[h])); - asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" - : "=v"(exp_sum[h]) - : "v"(exp_sum[h]), "v"(exp_sum[h])); - } - - if (laneid < 4) { - for (int h = 0; h < QHLOOP; h++) { - const int head_idx = 4 * h + lane4id; - shared_qk_max[warpid][head_idx] = qk_max[h]; - shared_exp_sum[warpid][head_idx] = exp_sum[h]; - } - } - } // warp within context - - __syncthreads(); - - const auto num_heads = gridDim.z * GQA_RATIO; - float* max_logits_ptr = - max_logits + seq_idx * num_heads * max_num_partitions + partition_idx; - float* exp_sums_ptr = - exp_sums + seq_idx * num_heads * max_num_partitions + partition_idx; - // calculate qk_max and exp_sums for partition - for (int h = 0; h < QHLOOP; h++) { - float global_qk_max = -FLT_MAX; - float warp_qk_max[NWARPS]; - const int head_idx = 4 * h + lane4id; - for (int w = 0; w < NWARPS; w++) { - warp_qk_max[w] = shared_qk_max[w][head_idx]; - global_qk_max = fmaxf(global_qk_max, warp_qk_max[w]); - } - float global_exp_sum = 0.0f; - for (int w = 0; w < NWARPS; w++) { - global_exp_sum += - shared_exp_sum[w][head_idx] * __expf(warp_qk_max[w] - global_qk_max); - } - if (head_idx < GQA_RATIO) { - max_logits_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = - global_qk_max; - exp_sums_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = - global_exp_sum; - } - const float global_inv_sum_scale = __fdividef(1.f, global_exp_sum + 1e-6f) * - __expf(qk_max[h] - global_qk_max); - d_out[h] *= global_inv_sum_scale; - } - constexpr bool LOGITS_RTZ_CONVERSION = false; - // logits[h] -> every 4 lanes hold 4 heads, each lane holds 4 tokens, there - // are 4x16 tokens across warp - _B16x4 logits[QHLOOP]; - for (int h = 0; h < QHLOOP; h++) { - if constexpr (LOGITS_RTZ_CONVERSION) { - // use rtz for faster performance with no perceivable accuracy loss - logits[h] = from_floatx4_rtz(d_out[h]); - } else { - logits[h] = from_floatx4(d_out[h]); - } - } - - if (warp_start_token_idx >= seq_len) { // warp out of context - for (int qh = 0; qh < QHLOOP; qh++) { - for (int vh = 0; vh < VHELOOP; vh++) { - vout_shared[qh][vh][laneid][warpid] = {0}; - } - } - } else { // warp in context - #define SV_mfma(x) \ - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ - Vlocal[vh][x] = convert_b8x8_custom(Vlocalb8[vh][x]); \ - } \ - for (int qh = 0; qh < QHLOOP; qh++) { \ - acc[qh] = gcn_mfma4x4x4_instr( \ - logits[qh], Vlocal[vh][x].xy[0], acc[qh]); \ - acc[qh] = gcn_mfma4x4x4_instr( \ - logits[qh], Vlocal[vh][x].xy[1], acc[qh]); \ - } - - for (int vh = 0; vh < VHELOOP; vh++) { - floatx4 acc[QHLOOP]; - for (int qh = 0; qh < QHLOOP; qh++) { - acc[qh] = {0}; - } - // SoftMax-V calculation - // logits -> token dimension is distributed across lanes - // Vlocal -> token dimension is depthwise within lane - // uses mfma instruction block broadcast for logits - SV_mfma(0); - SV_mfma(1); - SV_mfma(2); - SV_mfma(3); - SV_mfma(4); - SV_mfma(5); - SV_mfma(6); - SV_mfma(7); - - for (int qh = 0; qh < QHLOOP; qh++) { - if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { - // post mfma v scale for fp8 - acc[qh] *= *v_scale; - } - vout_shared[qh][vh][laneid][warpid] = from_floatx4(acc[qh]); - } - } - - #undef SV_mfma - } // warp in context - - __syncthreads(); - - // final write to tmp_out after vout accumulation - if (warpid == 0) { - _B16x4 vout[QHLOOP][VHELOOP]; - // iterate across heads - for (int qh = 0; qh < QHLOOP; qh++) { - // iterate over each v head elem (within head_size) - for (int vh = 0; vh < VHELOOP; vh++) { - vout[qh][vh] = {0}; - for (int w = 0; w < NWARPS; w++) { - vout[qh][vh] = - addx4(vout[qh][vh], vout_shared[qh][vh][laneid][w]); - } - } - } - - scalar_t* out_ptr = out + - seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - partition_idx * HEAD_SIZE; - const int out_num_partitions = max_num_partitions; - bit16_t* out_ptr_b16 = reinterpret_cast(out_ptr); - for (int qh = 0; qh < QHLOOP; qh++) { - for (int vh = 0; vh < VHELOOP; vh++) { - const int head_size_elem = vh * WARP_SIZE + laneid; - for (int i = 0; i < 4; i++) { - const int head_idx = 4 * qh + i; - if (head_idx < GQA_RATIO) { - out_ptr_b16[(wg_start_head_idx + head_idx) * out_num_partitions * - HEAD_SIZE + - head_size_elem] = vout[qh][vh][i]; - } - } - } - } - } // warpid == 0 -} - -// Grid: (num_heads, num_seqs). -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( - OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, - // max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { - const auto num_heads = gridDim.x; - const auto head_idx = blockIdx.x; - const auto seq_idx = blockIdx.y; - - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { - return; - } - - const int seq_len = seq_lens[seq_idx]; - const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); - const auto warpid = threadIdx.x / WARP_SIZE; - - __shared__ float shared_global_exp_sum; - // max num partitions supported is warp_size * NPAR_LOOPS - __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; - - if (warpid == 0) { - const float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - // valid partition is the last valid partition in case threadid > num - // partitions - int valid_partition[NPAR_LOOPS]; - float reg_max_logit[NPAR_LOOPS]; - const int last_valid_partition = num_partitions - 1; - - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const auto partition_no = i * WARP_SIZE + threadIdx.x; - valid_partition[i] = - (partition_no < num_partitions) ? partition_no : last_valid_partition; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; - } - float max_logit = reg_max_logit[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - max_logit = fmaxf(max_logit, reg_max_logit[i]); - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); - } - - const float* exp_sums_ptr = exp_sums + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - float rescaled_exp_sum[NPAR_LOOPS]; - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const auto partition_no = i * WARP_SIZE + threadIdx.x; - rescaled_exp_sum[i] *= (partition_no < num_partitions) - ? expf(reg_max_logit[i] - max_logit) - : 0.0f; - } - float global_exp_sum = rescaled_exp_sum[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - global_exp_sum += rescaled_exp_sum[i]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const auto partition_no = i * WARP_SIZE + threadIdx.x; - shared_exp_sums[partition_no] = rescaled_exp_sum[i]; - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - global_exp_sum += __shfl_xor(global_exp_sum, mask); - } - if (threadIdx.x == 0) { - shared_global_exp_sum = global_exp_sum; - } - } // warpid == 0 - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; - constexpr int MAX_NPAR = 64; - scalar_t tmps[MAX_NPAR]; - const float dzero = 0.0f; - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - tmps[j] = from_float(dzero); - } - const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; - const int num_partition_offset = (num_partitions)*HEAD_SIZE; - int idx = 0; - - constexpr int JCHUNK = 16; - - #pragma unroll - for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - __syncthreads(); - - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - } - } // num_partitions > JCHUNK - - // Aggregate tmp_out to out. - float acc = 0.0f; - #pragma unroll - for (int j = 0; j < JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK; j < 2 * JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - } - } - - for (int p = 1; p < NPAR_LOOPS; p++) { - if (num_partitions > p * MAX_NPAR) { - idx = 0; - #pragma unroll - for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; - } - } - } - - const float inv_global_exp_sum = - __fdividef(1.0f, shared_global_exp_sum + 1e-6f); - const float out_scale = - (fp8_out_scale_ptr != nullptr) ? 1.0f / (*fp8_out_scale_ptr) : 1.0f; - acc *= inv_global_exp_sum; - acc *= out_scale; - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + - static_cast(head_idx) * HEAD_SIZE; - if constexpr (std::is_same::value) { - out_ptr[threadIdx.x] = - __hip_cvt_float_to_fp8(acc, vllm::fp8::fp8_type::__default_saturation, - vllm::fp8::fp8_type::__default_interpret); - } else { - out_ptr[threadIdx.x] = from_float(acc); - } -} - -#elif defined(__GFX11__) - -using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; - -using bit16_t = uint16_t; -using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; -typedef bit16x4 _B16x4; - -using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; -union b16x8_u { - bit16x8 u16x8; - _B16x4 xy[2]; -}; -typedef b16x8_u _B16x8; - -using bit16x16 = - __attribute__((__vector_size__(16 * sizeof(uint16_t)))) uint16_t; -union b16x16_u { - bit16x16 u16x16; - _B16x8 xy[2]; -}; -typedef b16x16_u _B16x16; - -using _B8x8 = uint2; -using bit8_t = uint8_t; - -typedef struct _B8x16 { - _B8x8 xy[2]; -} _B8x16; - -template -__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x16& inpA, - const bit16x16& inpB, - const floatx8& inpC) { - if constexpr (std::is_same::value) { - return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(inpA, inpB, inpC); - } else if constexpr (std::is_same::value) { - return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(inpA, inpB, inpC); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ float to_float(const T& inp) { - if constexpr (std::is_same::value) { - return (float)inp; - } else if constexpr (std::is_same::value) { - return __bfloat162float(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ T from_float(const float& inp) { - if constexpr (std::is_same::value) { - return (_Float16)inp; - } else if constexpr (std::is_same::value) { - return __float2bfloat16(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { - if constexpr (std::is_same::value) { - union h2cvt { - __half2 h2[4]; - _B16x8 b16x8; - } u; - u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); - u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); - u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); - u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); - return u.b16x8; - } else if constexpr (std::is_same::value) { - union b2cvt { - __hip_bfloat162 b2[4]; - _B16x8 b16x8; - } u; - - u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); - u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); - u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); - u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); - - return u.b16x8; - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -// clang-format off -template -__global__ -__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - // clang-format on - constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 - const int warpid = threadIdx.x / WARP_SIZE; - const int laneid = threadIdx.x % WARP_SIZE; - const int lane2id = laneid % 2; - const int lane16id = laneid % 16; - const int rowid = laneid / 16; - - const int seq_idx = blockIdx.x; - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { - return; - } - - const int partition_idx = blockIdx.y; - - constexpr int T_PAR_SIZE = 256; // token partition size set to 256 - - const int max_num_partitions = gridDim.y; - - const int seq_len = seq_lens[seq_idx]; // length of a seq - - const int partition_start_token_idx = partition_idx * T_PAR_SIZE; - // exit if partition is out of context for seq - if (partition_start_token_idx >= seq_len) { - return; - } - - constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); - - __shared__ float shared_qk_max[NWARPS][16 + 1]; - __shared__ float shared_exp_sum[NWARPS][16 + 1]; - // shared_logits is used for multiple purposes - __shared__ _B16x16 shared_logits[NWARPS][2][16][2]; - - // for QK wmma16x16, layout is QHead/Tokenx16 across every 16 lanes, - // 32 Bytes HeadElements in each lane, 2x16B HeadElements across a row of warp - constexpr int ROWS_PER_WARP = - WARP_SIZE / 16 / 2; // rows refers to 16 lanes; refer dpp terminology - constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = - 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types - constexpr int QKHE_PER_FETCH = - CONTIGUOUS_KV_ELEMS_16B_LOAD * - ROWS_PER_WARP; // each fetch across a warp fetches these many elements - constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across - // warp - - _B16x16 Qlocal[QKHELOOP / 2]; // note that 16 contiguous elements of Q should - // be fetched per lane for 16 bit cache types - - constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); - - constexpr int TOKENS_PER_WARP = - T_PAR_SIZE / - NWARPS; // sub partition of tokens per warp for qk calculation - constexpr int TLOOP = - TOKENS_PER_WARP / - 16; // each wmma16x16x16 instruction processes 16 tokens - - _B16x16 Klocal[TLOOP] - [QKHELOOP / 2]; // can be interpreted as B8x16 for 8 bit types - - const int wg_start_head_idx = blockIdx.z * GQA_RATIO; - const int wg_start_kv_head_idx = blockIdx.z; - const int total_num_heads = gridDim.z * GQA_RATIO; - - // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps - // each wmma takes QH16xT16x16HE across warp - // repeat wmma across QKHELOOP dimension - // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens - // across 2 rows x 8 tokens per lane - - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - - if (GQA_RATIO == 1) { - const int local_qhead_idx = lane16id % GQA_RATIO; - const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; - const scalar_t* q_ptr = - q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; - if (lane16id < GQA_RATIO) { - #pragma unroll - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { - const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH * 2; - const _B16x16* q_fetch_ptr_32B = - reinterpret_cast(q_fetch_ptr); - Qlocal[qkhe_depth] = *q_fetch_ptr_32B; - } - } - } else { - // fetch Q in shared across warps and then write to registers - const int local_qhead_idx = 2 * warpid + rowid; - const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; - const scalar_t* q_ptr = - q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; - - const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; - if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { - const scalar_t* q_fetch_ptr = q_ptr + qhead_element; - const _B16x8* q_fetch_ptr_16B = - reinterpret_cast(q_fetch_ptr); - _B16x8 tmp = *q_fetch_ptr_16B; - - const int offset1 = - lane16id / - 2; // 16 contiguous chunks of head elems are spread across 8x2lanes - shared_logits[offset1][lane2id][local_qhead_idx][0].xy[0] = tmp; - } - - __syncthreads(); - - #pragma unroll - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { - Qlocal[qkhe_depth].xy[0] = - shared_logits[qkhe_depth][0][lane16id % GQA_RATIO][0].xy[0]; - Qlocal[qkhe_depth].xy[1] = - shared_logits[qkhe_depth][1][lane16id % GQA_RATIO][0].xy[0]; - } - } - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int last_seq_block = num_seq_blocks - 1; - - const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; - - int kphysical_block_number[TLOOP]; - - // fetch k physical block numbers - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; - const int kblock_idx = (kglobal_token_idx < seq_len) - ? kglobal_token_idx / BLOCK_SIZE - : last_seq_block; - kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; - } - - constexpr int KX = 16 / sizeof(cache_t); - const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; - - const int row_head_elem = 0; - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int64_t kblock_number = - static_cast(kphysical_block_number[token_depth]); - const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; - const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; - - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; - const int offset1 = head_elem / KX; - const int offset2 = head_elem % KX; - const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; - const _B16x8* k_fetch_ptr_16B = - reinterpret_cast(k_fetch_ptr); - Klocal[token_depth][qkhe_depth / 2].xy[qkhe_depth % 2] = *k_fetch_ptr_16B; - } - } - - constexpr int VTOKENS_PER_LANE = - TOKENS_PER_WARP / ROWS_PER_WARP; // 32/1 = 32 vtokens per lane - constexpr int VBLOCKS_PER_LANE = 2; // assumes block size >=16 - constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps - constexpr int VTLANELOOP = DIVIDE_ROUND_UP( - VTOKENS_PER_LANE, - CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes - // minimum block size is 16 - constexpr int VHELOOP = DIVIDE_ROUND_UP( - (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each - // wmma instr works on 16 head elements - - int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; - - // fetch v physical block numbers - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; - vblock_depth++) { - const int vlocal_token_idx = - vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + - vblock_depth * BLOCK_SIZE; - const int vglobal_token_idx = - partition_start_token_idx + vlocal_token_idx; - const int vblock_idx = (vglobal_token_idx < seq_len) - ? vglobal_token_idx / BLOCK_SIZE - : last_seq_block; - vphysical_block_number[vtoken_depth][vblock_depth] = - block_table_seq[vblock_idx]; - } - } - - _B16x16 Vlocal[VTLOOP][VHELOOP] - [VTLANELOOP / 2]; // this can be interpreted as B8x16 too - - const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; - // v fetches are 16head elems across lanes x (16x2) tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; - const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - const int64_t vblock_number = static_cast( - vphysical_block_number[vtoken_depth] - [vfetch_depth / VBLOCKS_PER_LANE]); - const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); - - const cache_t* v_fetch_ptr = - v_ptr3 + - (vfetch_depth % VBLOCKS_PER_LANE) * CONTIGUOUS_KV_ELEMS_16B_LOAD; - const _B16x8* v_fetch_ptr_16B = - reinterpret_cast(v_fetch_ptr); - Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2].xy[vfetch_depth % 2] = - *v_fetch_ptr_16B; - } - } - } - - floatx8 dout[TLOOP]; - // qk wmma - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - dout[token_depth] = {0}; - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { - dout[token_depth] = gcn_wmma16x16x16_instr( - Klocal[token_depth][qkhe_depth].u16x16, Qlocal[qkhe_depth].u16x16, - dout[token_depth]); - } - dout[token_depth] *= scale; - } - - // calculate qk_max and exp_sum per warp and write to shared memory - float qk_max = -FLT_MAX; - float exp_sum = 0.0f; - const int qkout_token_idx = - partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid; - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 8; i++) { - const float tmp = - (local_token_idx + 2 * i < seq_len) ? dout[token_depth][i] : -FLT_MAX; - qk_max = fmaxf(qk_max, tmp); - } - } - - qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 8; i++) { - const float tmp = (local_token_idx + 2 * i < seq_len) - ? __expf(dout[token_depth][i] - qk_max) - : 0.0f; - dout[token_depth][i] = tmp; - exp_sum += tmp; - } - } - - exp_sum += __shfl_xor(exp_sum, 16); - - __syncthreads(); - - if (laneid < 16) { - shared_qk_max[warpid][lane16id] = qk_max; - shared_exp_sum[warpid][lane16id] = exp_sum; - } - - __syncthreads(); - - // calculate partition qk_max and exp_sum - float partition_qk_max = -FLT_MAX; - float warp_qk_max_exp[NWARPS]; - float partition_exp_sum = 0.0f; - - #pragma unroll - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; - partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); - } - - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); - partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; - } - - const float inv_sum_scale = - __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; - - __syncthreads(); - - // write logits to shared mem - #pragma unroll - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - dout[token_depth] *= inv_sum_scale; - shared_logits[warpid][token_depth][lane16id][0].xy[rowid] = - from_floatx8(dout[token_depth]); - } - __syncthreads(); - - _B16x8 swp_buf[TLOOP][2]; - #pragma unroll - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - swp_buf[token_depth][0] = - shared_logits[warpid][token_depth][lane16id][0].xy[0]; - swp_buf[token_depth][1] = - shared_logits[warpid][token_depth][lane16id][0].xy[1]; - } - - #pragma unroll - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - #pragma unroll - for (int i = 0; i < 8; i++) { - shared_logits[warpid][token_depth][lane16id][0].xy[rowid].u16x8[i] = - swp_buf[token_depth][i % 2].u16x8[4 * rowid + (i / 2)]; - } - } - - // write out partition max_logits and exp_sum - if (threadIdx.x < GQA_RATIO) { - const int qhead_idx = lane16id; - const int offset = seq_idx * total_num_heads * max_num_partitions + - (wg_start_head_idx + qhead_idx) * max_num_partitions + - partition_idx; - max_logits[offset] = partition_qk_max; - exp_sums[offset] = partition_exp_sum; - } - - __syncthreads(); - - _B16x8 outelems[VHELOOP]; - // Softmax V wmma - // v layout: 16he across lanes x (16x2) tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - floatx8 tmp_out = {0}; - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP / 2; - vfetch_depth++) { - const int offset = vfetch_depth; - // if output format is 16 qheads across 16 lanes, 16 head elems spread - // across rows - tmp_out = gcn_wmma16x16x16_instr( - Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x16, - shared_logits[vtoken_depth][offset][lane16id][0].u16x16, tmp_out); - } - } - outelems[vhe_depth] = from_floatx8(tmp_out); - } - - __syncthreads(); - - #pragma unroll - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid] = - outelems[vhe_depth]; // lane16 id head dimension; rowid head element - // dimension - } - - __syncthreads(); - - #pragma unroll - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - swp_buf[vhe_depth][0] = shared_logits[warpid][vhe_depth][lane16id][0].xy[0]; - swp_buf[vhe_depth][1] = shared_logits[warpid][vhe_depth][lane16id][0].xy[1]; - } - - #pragma unroll - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - #pragma unroll - for (int i = 0; i < 8; i++) { - shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid].u16x8[i] = - swp_buf[vhe_depth][i % 2].u16x8[4 * rowid + (i / 2)]; - } - } - - __syncthreads(); - - // write to tmp_out with coalesced writes after reading from shared mem - if (warpid == 0) { - _B16x8 vout[GQA_RATIO2]; - // each lane writes out 16Bytes of tmp_out along head elem dimension - const int head_elem_idx = lane16id * 8; - if (head_elem_idx < HEAD_SIZE) { - for (int h = 0; h < GQA_RATIO2; h++) { - const int local_head_idx = 2 * h + rowid; - const int offset1 = (head_elem_idx / 16) % NWARPS; - const int offset2 = head_elem_idx / 16 / NWARPS; - const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row - vout[h] = - shared_logits[offset1][offset2][local_head_idx][0].xy[offset3]; - } - - const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; - scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + - partition_idx * HEAD_SIZE; - for (int h = 0; h < GQA_RATIO2; h++) { - const int local_head_idx = 2 * h + rowid; - if (local_head_idx < GQA_RATIO) { - const int out_head_idx = wg_start_head_idx + local_head_idx; - scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; - scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; - _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); - *out_ptr_B16x8 = vout[h]; - } - } - } - } -} - -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - UNREACHABLE_CODE -} - -// Grid: (num_heads, num_seqs). -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( - OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, - // max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { - const auto num_heads = gridDim.x; - const auto head_idx = blockIdx.x; - const auto seq_idx = blockIdx.y; - - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { - return; - } - - const int seq_len = seq_lens[seq_idx]; - const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); - const int warpid = threadIdx.x / WARP_SIZE; - - __shared__ float shared_global_exp_sum; - // max num partitions supported is warp_size * NPAR_LOOPS - __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; - - if (warpid == 0) { - const float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - // valid partition is the last valid partition in case threadid > num - // partitions - int valid_partition[NPAR_LOOPS]; - float reg_max_logit[NPAR_LOOPS]; - const int last_valid_partition = num_partitions - 1; - - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - valid_partition[i] = - (partition_no < num_partitions) ? partition_no : last_valid_partition; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; - } - float max_logit = reg_max_logit[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - max_logit = fmaxf(max_logit, reg_max_logit[i]); - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); - } - - const float* exp_sums_ptr = exp_sums + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - float rescaled_exp_sum[NPAR_LOOPS]; - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - rescaled_exp_sum[i] *= (partition_no < num_partitions) - ? expf(reg_max_logit[i] - max_logit) - : 0.0f; - } - float global_exp_sum = rescaled_exp_sum[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - global_exp_sum += rescaled_exp_sum[i]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - shared_exp_sums[partition_no] = rescaled_exp_sum[i]; - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - global_exp_sum += __shfl_xor(global_exp_sum, mask); - } - if (threadIdx.x == 0) { - shared_global_exp_sum = global_exp_sum; - } - } // warpid == 0 - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; - constexpr int MAX_NPAR = 32; - scalar_t tmps[MAX_NPAR]; - const float dzero = 0.0f; - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - tmps[j] = from_float(dzero); - } - const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; - const int num_partition_offset = (num_partitions)*HEAD_SIZE; - int idx = 0; - - constexpr int JCHUNK = 16; - - #pragma unroll - for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - __syncthreads(); - - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - } - } // num_partitions > JCHUNK - - // Aggregate tmp_out to out. - float acc = 0.0f; - #pragma unroll - for (int j = 0; j < JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK; j < 2 * JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - } - } - - for (int p = 1; p < NPAR_LOOPS; p++) { - if (num_partitions > p * MAX_NPAR) { - idx = 0; - #pragma unroll - for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; - } - } - } - - const float inv_global_exp_sum = - __fdividef(1.0f, shared_global_exp_sum + 1e-6f); - acc *= inv_global_exp_sum; - - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + - static_cast(head_idx) * HEAD_SIZE; - out_ptr[threadIdx.x] = from_float(acc); -} - -#elif defined(__GFX12__) - -using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; - -using bit16_t = uint16_t; -using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; -typedef bit16x4 _B16x4; - -using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; -union b16x8_u { - bit16x8 u16x8; - _B16x4 xy[2]; -}; -typedef b16x8_u _B16x8; - -using _B8x8 = uint2; -using bit8_t = uint8_t; - -typedef struct _B8x16 { - _B8x8 xy[2]; -} _B8x16; - -template -__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, - const bit16x8& inpB, - const floatx8& inpC) { - if constexpr (std::is_same::value) { - return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); - } else if constexpr (std::is_same::value) { - return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(inpA, inpB, inpC); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ float to_float(const T& inp) { - if constexpr (std::is_same::value) { - return (float)inp; - } else if constexpr (std::is_same::value) { - return __bfloat162float(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ float to_float_b16(const bit16_t& inp) { - union tmpcvt { - bit16_t u; - _Float16 f; - __hip_bfloat16 b; - } t16; - t16.u = inp; - if constexpr (std::is_same::value) { - return (float)t16.f; - } else if constexpr (std::is_same::value) { - return __bfloat162float(t16.b); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ T from_float(const float& inp) { - if constexpr (std::is_same::value) { - return (_Float16)inp; - } else if constexpr (std::is_same::value) { - return __float2bfloat16(inp); - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -template -__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { - if constexpr (std::is_same::value) { - union h2cvt { - __half2 h2[4]; - _B16x8 b16x8; - } u; - u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); - u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); - u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); - u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); - return u.b16x8; - } else if constexpr (std::is_same::value) { - union b2cvt { - __hip_bfloat162 b2[4]; - _B16x8 b16x8; - } u; - - u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); - u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); - u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); - u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); - - return u.b16x8; - } else { - static_assert(false, "unsupported 16b dtype"); - } -} - -// clang-format off -template -__global__ -__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - // clang-format on - constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 - const int warpid = threadIdx.x / WARP_SIZE; - const int laneid = threadIdx.x % WARP_SIZE; - const int lane2id = laneid % 2; - const int lane16id = laneid % 16; - const int rowid = laneid / 16; - - const int seq_idx = blockIdx.x; - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { - return; - } - const int partition_idx = blockIdx.y; - - constexpr int T_PAR_SIZE = 256; // token partition size set to 256 - - const int max_num_partitions = gridDim.y; - - const int seq_len = seq_lens[seq_idx]; // length of a seq - - const int partition_start_token_idx = partition_idx * T_PAR_SIZE; - // exit if partition is out of context for seq - if (partition_start_token_idx >= seq_len) { - return; - } - - constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); - - __shared__ float shared_qk_max[NWARPS][16 + 1]; - __shared__ float shared_exp_sum[NWARPS][16 + 1]; - // shared_logits is used for multiple purposes - __shared__ _B16x8 shared_logits[NWARPS][2][16][2]; - - // for QK wmma16x16_gfx12, layout is QHead/Tokenx16 across every 16 lanes, - // 16 Bytes HeadElements in each lane, 2x16B HeadElements across 2 rows of - // warp - constexpr int ROWS_PER_WARP = - WARP_SIZE / 16; // rows refers to 16 lanes; refer dpp terminology - constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = - 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types - constexpr int QKHE_PER_FETCH = - CONTIGUOUS_KV_ELEMS_16B_LOAD * - ROWS_PER_WARP; // each fetch across a warp fetches these many elements - constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across - // warp - - _B16x8 Qlocal[QKHELOOP]; // note that 16 contiguous elements of Q should - // be fetched per lane for 16 bit cache types - - constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); - - constexpr int TOKENS_PER_WARP = - T_PAR_SIZE / - NWARPS; // sub partition of tokens per warp for qk calculation - constexpr int TLOOP = - TOKENS_PER_WARP / - 16; // each wmma16x16x16 instruction processes 16 tokens - - _B16x8 Klocal[TLOOP] - [QKHELOOP]; // can be interpreted as B8x16 for 8 bit types - - const int wg_start_head_idx = blockIdx.z * GQA_RATIO; - const int wg_start_kv_head_idx = blockIdx.z; - const int total_num_heads = gridDim.z * GQA_RATIO; - - // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps - // each wmma takes QH16xT16x16HE across warp - // repeat wmma across QKHELOOP dimension - // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens - // across 2 rows x 8 tokens per lane - - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - - if (GQA_RATIO == 1) { - const int local_qhead_idx = lane16id % GQA_RATIO; - const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; - const scalar_t* q_ptr = q + query_start_off * q_stride + - global_qhead_idx * HEAD_SIZE + - rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; - if (lane16id < GQA_RATIO) { - #pragma unroll - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH; - const _B16x8* q_fetch_ptr_16B = - reinterpret_cast(q_fetch_ptr); - Qlocal[qkhe_depth] = *q_fetch_ptr_16B; - } - } - } else { - // fetch Q in shared across warps and then write to registers - const int local_qhead_idx = 2 * warpid + rowid; - const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; - const scalar_t* q_ptr = - q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; - - const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; - if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { - const scalar_t* q_fetch_ptr = q_ptr + qhead_element; - const _B16x8* q_fetch_ptr_16B = - reinterpret_cast(q_fetch_ptr); - _B16x8 tmp = *q_fetch_ptr_16B; - - const int offset1 = - lane16id / - 2; // 16 contiguous chunks of head elems are spread across 8x2lanes - shared_logits[offset1][lane2id][local_qhead_idx][0] = tmp; - } - - __syncthreads(); - - #pragma unroll - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - Qlocal[qkhe_depth] = - shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO][0]; - } - } - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int last_seq_block = num_seq_blocks - 1; - - const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; - - int kphysical_block_number[TLOOP]; - - // fetch k physical block numbers - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; - const int kblock_idx = (kglobal_token_idx < seq_len) - ? kglobal_token_idx / BLOCK_SIZE - : last_seq_block; - kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; - } - - constexpr int KX = 16 / sizeof(cache_t); - const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; - - const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int64_t kblock_number = - static_cast(kphysical_block_number[token_depth]); - const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; - const int klocal_token_idx = - TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; - const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; - const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; - - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; - const int offset1 = head_elem / KX; - const int offset2 = head_elem % KX; - const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; - const _B16x8* k_fetch_ptr_16B = - reinterpret_cast(k_fetch_ptr); - Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; - } - } - - constexpr int VTOKENS_PER_LANE = - TOKENS_PER_WARP / ROWS_PER_WARP; // 32/2 = 16 vtokens per lane - constexpr int VBLOCKS_PER_LANE = 1; // assumes block size >=16 - constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps - constexpr int VTLANELOOP = DIVIDE_ROUND_UP( - VTOKENS_PER_LANE, - CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes - // minimum block size is 16 - constexpr int VHELOOP = DIVIDE_ROUND_UP( - (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each - // wmma instr works on 16 head elements - - int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; - - // fetch v physical block numbers - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; - vblock_depth++) { - const int vlocal_token_idx = - vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + - rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; - const int vglobal_token_idx = - partition_start_token_idx + vlocal_token_idx; - const int vblock_idx = (vglobal_token_idx < seq_len) - ? vglobal_token_idx / BLOCK_SIZE - : last_seq_block; - vphysical_block_number[vtoken_depth][vblock_depth] = - block_table_seq[vblock_idx]; - } - } - - _B16x8 Vlocal[VTLOOP][VHELOOP] - [VTLANELOOP]; // this can be interpreted as B8x16 too - - const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + - ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); - - // v fetches are 16head elems across lanes x 16 tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; - const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - const int vblock_depth = 0; - const int64_t vblock_number = static_cast( - vphysical_block_number[vtoken_depth][vblock_depth]); - const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); - - const cache_t* v_fetch_ptr = - v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; - const _B16x8* v_fetch_ptr_16B = - reinterpret_cast(v_fetch_ptr); - Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; - } - } - } - - floatx8 dout[TLOOP]; - // qk wmma - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - dout[token_depth] = {0}; - for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { - dout[token_depth] = gcn_wmma16x16x16_instr( - Klocal[token_depth][qkhe_depth].u16x8, Qlocal[qkhe_depth].u16x8, - dout[token_depth]); - } - dout[token_depth] *= scale; - } - - // calculate qk_max and exp_sum per warp and write to shared memory - float qk_max = -FLT_MAX; - float exp_sum = 0.0f; - const int qkout_token_idx = - partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 8; - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 8; i++) { - const float tmp = - (local_token_idx + i < seq_len) ? dout[token_depth][i] : -FLT_MAX; - qk_max = fmaxf(qk_max, tmp); - } - } - - qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); - - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - const int local_token_idx = qkout_token_idx + token_depth * 16; - for (int i = 0; i < 8; i++) { - const float tmp = (local_token_idx + i < seq_len) - ? __expf(dout[token_depth][i] - qk_max) - : 0.0f; - dout[token_depth][i] = tmp; - exp_sum += tmp; - } - } - - exp_sum += __shfl_xor(exp_sum, 16); - - __syncthreads(); - - if (laneid < 16) { - shared_qk_max[warpid][lane16id] = qk_max; - shared_exp_sum[warpid][lane16id] = exp_sum; - } - - __syncthreads(); - - // calculate partition qk_max and exp_sum - float partition_qk_max = -FLT_MAX; - float warp_qk_max_exp[NWARPS]; - float partition_exp_sum = 0.0f; - - #pragma unroll - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; - partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); - } - - for (int w = 0; w < NWARPS; w++) { - warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); - partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; - } - - const float inv_sum_scale = - __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; - - __syncthreads(); - - // write logits to shared mem - #pragma unroll - for (int token_depth = 0; token_depth < TLOOP; token_depth++) { - dout[token_depth] *= inv_sum_scale; - shared_logits[warpid][token_depth][lane16id][rowid] = - from_floatx8(dout[token_depth]); - } - - // write out partition max_logits and exp_sum - if (threadIdx.x < GQA_RATIO) { - const int qhead_idx = lane16id; - const int offset = seq_idx * total_num_heads * max_num_partitions + - (wg_start_head_idx + qhead_idx) * max_num_partitions + - partition_idx; - max_logits[offset] = partition_qk_max; - exp_sums[offset] = partition_exp_sum; - } - - __syncthreads(); - - _B16x8 outelems[VHELOOP]; - // Softmax V wmma - // v layout: 16he across lanes x 16 tokens per lane - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - floatx8 tmp_out = {0}; - - for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { - for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { - const int offset = rowid * VTLANELOOP + vfetch_depth; - const int offset1 = offset % ROWS_PER_WARP; - const int offset2 = offset / ROWS_PER_WARP; - // if output format is 16 qheads across 16 lanes, 16 head elems spread - // across rows - tmp_out = gcn_wmma16x16x16_instr( - Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x8, - shared_logits[vtoken_depth][offset2][lane16id][offset1].u16x8, - tmp_out); - } - } - outelems[vhe_depth] = from_floatx8(tmp_out); - } - - __syncthreads(); - - #pragma unroll - for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { - shared_logits[warpid][vhe_depth][lane16id][rowid] = - outelems[vhe_depth]; // lane16 id head dimension; rowid head element - // dimension - } - - __syncthreads(); - - // write to tmp_out with coalesced writes after reading from shared mem - if (warpid == 0) { - _B16x8 vout[GQA_RATIO2]; - // each lane writes out 16Bytes of tmp_out along head elem dimension - const int head_elem_idx = lane16id * 8; - if (head_elem_idx < HEAD_SIZE) { - for (int h = 0; h < GQA_RATIO2; h++) { - const int local_head_idx = 2 * h + rowid; - const int offset1 = (head_elem_idx / 16) % NWARPS; - const int offset2 = head_elem_idx / 16 / NWARPS; - const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row - vout[h] = shared_logits[offset1][offset2][local_head_idx][offset3]; - } - - const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; - scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + - partition_idx * HEAD_SIZE; - for (int h = 0; h < GQA_RATIO2; h++) { - const int local_head_idx = 2 * h + rowid; - if (local_head_idx < GQA_RATIO) { - const int out_head_idx = wg_start_head_idx + local_head_idx; - scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; - scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; - _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); - *out_ptr_B16x8 = vout[h]; - } - } - } - } -} - -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - UNREACHABLE_CODE -} - -// Grid: (num_heads, num_seqs). -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( - OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, - // max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { - const auto num_heads = gridDim.x; - const auto head_idx = blockIdx.x; - const auto seq_idx = blockIdx.y; - - // NOTE queries with sequence len > 1 are prefills and taken care by another - // kernel. - if (query_start_loc_ptr != nullptr && - (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { - return; - } - - const int seq_len = seq_lens[seq_idx]; - const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); - const int warpid = threadIdx.x / WARP_SIZE; - - __shared__ float shared_global_exp_sum; - // max num partitions supported is warp_size * NPAR_LOOPS - __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; - - if (warpid == 0) { - const float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - // valid partition is the last valid partition in case threadid > num - // partitions - int valid_partition[NPAR_LOOPS]; - float reg_max_logit[NPAR_LOOPS]; - const int last_valid_partition = num_partitions - 1; - - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - valid_partition[i] = - (partition_no < num_partitions) ? partition_no : last_valid_partition; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; - } - float max_logit = reg_max_logit[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - max_logit = fmaxf(max_logit, reg_max_logit[i]); - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); - } - - const float* exp_sums_ptr = exp_sums + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - - float rescaled_exp_sum[NPAR_LOOPS]; - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - rescaled_exp_sum[i] *= (partition_no < num_partitions) - ? expf(reg_max_logit[i] - max_logit) - : 0.0f; - } - float global_exp_sum = rescaled_exp_sum[0]; - #pragma unroll - for (int i = 1; i < NPAR_LOOPS; i++) { - global_exp_sum += rescaled_exp_sum[i]; - } - #pragma unroll - for (int i = 0; i < NPAR_LOOPS; i++) { - const int partition_no = i * WARP_SIZE + threadIdx.x; - shared_exp_sums[partition_no] = rescaled_exp_sum[i]; - } - - #pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - global_exp_sum += __shfl_xor(global_exp_sum, mask); - } - if (threadIdx.x == 0) { - shared_global_exp_sum = global_exp_sum; - } - } // warpid == 0 - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; - constexpr int MAX_NPAR = 32; - scalar_t tmps[MAX_NPAR]; - const float dzero = 0.0f; - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - tmps[j] = from_float(dzero); - } - const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; - const int num_partition_offset = (num_partitions)*HEAD_SIZE; - int idx = 0; - - constexpr int JCHUNK = 16; - - #pragma unroll - for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - __syncthreads(); - - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - } - } // num_partitions > JCHUNK - - // Aggregate tmp_out to out. - float acc = 0.0f; - #pragma unroll - for (int j = 0; j < JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > JCHUNK) { - #pragma unroll - for (int j = JCHUNK; j < 2 * JCHUNK; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - if (num_partitions > 2 * JCHUNK) { - #pragma unroll - for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j]; - } - } - } - - for (int p = 1; p < NPAR_LOOPS; p++) { - if (num_partitions > p * MAX_NPAR) { - idx = 0; - #pragma unroll - for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; - j += HEAD_SIZE) { - // lastj is last valid partition - const int lastj_offset = - (j < num_partition_offset) ? j : last_partition_offset; - tmps[idx] = tmp_out_ptr[lastj_offset]; - idx++; - } - - #pragma unroll - for (int j = 0; j < MAX_NPAR; j++) { - acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; - } - } - } - - const float inv_global_exp_sum = - __fdividef(1.0f, shared_global_exp_sum + 1e-6f); - acc *= inv_global_exp_sum; - - const int64_t query_start_off = static_cast( - query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); - OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + - static_cast(head_idx) * HEAD_SIZE; - out_ptr[threadIdx.x] = from_float(acc); -} - -#else - -// clang-format off -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] - const int num_kv_heads, - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, - const int kv_block_stride, - const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - UNREACHABLE_CODE -} - -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] - const int num_kv_heads, - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, - const int kv_block_stride, - const int kv_head_stride, - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] - OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] - int max_ctx_blocks, const float* k_scale, const float* v_scale) { - UNREACHABLE_CODE -} - -// Grid: (num_heads, num_seqs). -template -__global__ -__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( - OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int* __restrict__ query_start_loc_ptr, // [num_seqs] - const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { - UNREACHABLE_CODE -} -// clang-format on - -#endif - -#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \ - paged_attention_ll4mi_QKV_mfma16_kernel \ - <<>>( \ - query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ - block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ - max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ - kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ - max_ctx_blocks, k_scale_ptr, v_scale_ptr); - -#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \ - paged_attention_ll4mi_QKV_mfma4_kernel \ - <<>>( \ - query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ - block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ - max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ - kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ - max_ctx_blocks, k_scale_ptr, v_scale_ptr); - -#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \ - paged_attention_ll4mi_reduce_kernel \ - <<>>( \ - out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \ - query_start_loc_ptr, max_num_partitions, fp8_out_scale_ptr); - -template -void paged_attention_custom_launcher( +// Per-dtype dispatch shards (defined in attention/dispatch_auto.cu and +// attention/dispatch_fp8.cu). +void paged_attention_dispatch_auto( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& value_cache, int64_t num_kv_heads, double scale, torch::Tensor& block_tables, torch::Tensor& seq_lens, - const std::optional& query_start_loc, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale, const std::optional& fp8_out_scale) { - int num_seqs = block_tables.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: query start location is optional for V0 decode should not be used. - // If batch contains mix of prefills and decode, prefills should be skipped. - const int* query_start_loc_ptr = - query_start_loc - ? reinterpret_cast(query_start_loc.value().data_ptr()) - : nullptr; - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; + const std::optional& query_start_loc, int64_t block_size, + int64_t max_seq_len, const std::optional& alibi_slopes, + torch::Tensor& k_scale, torch::Tensor& v_scale, + const std::optional& fp8_out_scale, bool is_navi); - float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); - float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); - T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - // NOTE: fp8_out_scale is optional. - const auto fp8_out_scale_ptr = - fp8_out_scale - ? static_cast(fp8_out_scale.value().data_ptr()) - : nullptr; - OUTT* out_ptr = reinterpret_cast(out.data_ptr()); - - const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); - - // partition size is fixed at 256 since both mfma4 and mfma16 kernels support - // it mfma4 kernel also supports partition size 512 - constexpr int PARTITION_SIZE = 256; - const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); - const int gqa_ratio = num_heads / num_kv_heads; - assert(num_heads % num_kv_heads == 0); - assert(head_size == HEAD_SIZE); - - constexpr int NTHR = 256; - dim3 grid(num_seqs, max_num_partitions, num_kv_heads); - dim3 block(NTHR); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - // mfma4 kernel is faster than mfma16 for gqa_ratio <= 4 - switch (gqa_ratio) { - case 1: - LAUNCH_CUSTOM_ATTENTION_MFMA4(1); - break; - case 2: - LAUNCH_CUSTOM_ATTENTION_MFMA4(2); - break; - case 3: - LAUNCH_CUSTOM_ATTENTION_MFMA4(3); - break; - case 4: - LAUNCH_CUSTOM_ATTENTION_MFMA4(4); - break; - case 5: - LAUNCH_CUSTOM_ATTENTION_MFMA16(5); - break; - case 6: - LAUNCH_CUSTOM_ATTENTION_MFMA16(6); - break; - case 7: - LAUNCH_CUSTOM_ATTENTION_MFMA16(7); - break; - case 8: - LAUNCH_CUSTOM_ATTENTION_MFMA16(8); - break; - case 9: - LAUNCH_CUSTOM_ATTENTION_MFMA16(9); - break; - case 10: - LAUNCH_CUSTOM_ATTENTION_MFMA16(10); - break; - case 11: - LAUNCH_CUSTOM_ATTENTION_MFMA16(11); - break; - case 12: - LAUNCH_CUSTOM_ATTENTION_MFMA16(12); - break; - case 13: - LAUNCH_CUSTOM_ATTENTION_MFMA16(13); - break; - case 14: - LAUNCH_CUSTOM_ATTENTION_MFMA16(14); - break; - case 15: - LAUNCH_CUSTOM_ATTENTION_MFMA16(15); - break; - case 16: - LAUNCH_CUSTOM_ATTENTION_MFMA16(16); - break; - default: - TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); - break; - } - - dim3 reduce_grid(num_heads, num_seqs); - dim3 reduce_block(head_size); - const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, WARP_SIZE); - // reduction kernel supports upto 8 NPAR_loops * 64 (warp_size) * 256 - // (partition size) = 128K context length - switch (npar_loops) { - case 1: - LAUNCH_CUSTOM_REDUCTION(1); - break; - case 2: - LAUNCH_CUSTOM_REDUCTION(2); - break; - case 3: - LAUNCH_CUSTOM_REDUCTION(3); - break; - case 4: - LAUNCH_CUSTOM_REDUCTION(4); - break; - case 5: - LAUNCH_CUSTOM_REDUCTION(5); - break; - case 6: - LAUNCH_CUSTOM_REDUCTION(6); - break; - case 7: - LAUNCH_CUSTOM_REDUCTION(7); - break; - case 8: - LAUNCH_CUSTOM_REDUCTION(8); - break; - default: - TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); - break; - } -} - -template -void paged_attention_custom_launcher_navi( +void paged_attention_dispatch_fp8( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, - torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& value_cache, int64_t num_kv_heads, double scale, torch::Tensor& block_tables, torch::Tensor& seq_lens, - const std::optional& query_start_loc, int max_seq_len, - const std::optional& alibi_slopes, torch::Tensor& k_scale, - torch::Tensor& v_scale) { - int num_seqs = block_tables.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: query start location is optional for V0 decode should not be used. - // If batch contains mix of prefills and decode, prefills should be skipped. - const int* query_start_loc_ptr = - query_start_loc - ? reinterpret_cast(query_start_loc.value().data_ptr()) - : nullptr; - - // NOTE: Navi does not support alibi_slopes. - const float* alibi_slopes_ptr = nullptr; - - float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); - float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); - T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.data_ptr(); - int* seq_lens_ptr = seq_lens.data_ptr(); - - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - // NOTE: Navi does not support fp8. - const auto fp8_out_scale_ptr = nullptr; - OUTT* out_ptr = reinterpret_cast(out.data_ptr()); - - const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); - - constexpr int PARTITION_SIZE = 256; - const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); - const int gqa_ratio = num_heads / num_kv_heads; - assert(num_heads % num_kv_heads == 0); - assert(head_size == HEAD_SIZE); - - constexpr int NTHR = 256; - dim3 grid(num_seqs, max_num_partitions, num_kv_heads); - dim3 block(NTHR); - const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - switch (gqa_ratio) { - case 1: - LAUNCH_CUSTOM_ATTENTION_MFMA16(1); - break; - case 2: - LAUNCH_CUSTOM_ATTENTION_MFMA16(2); - break; - case 3: - LAUNCH_CUSTOM_ATTENTION_MFMA16(3); - break; - case 4: - LAUNCH_CUSTOM_ATTENTION_MFMA16(4); - break; - case 5: - LAUNCH_CUSTOM_ATTENTION_MFMA16(5); - break; - case 6: - LAUNCH_CUSTOM_ATTENTION_MFMA16(6); - break; - case 7: - LAUNCH_CUSTOM_ATTENTION_MFMA16(7); - break; - case 8: - LAUNCH_CUSTOM_ATTENTION_MFMA16(8); - break; - case 9: - LAUNCH_CUSTOM_ATTENTION_MFMA16(9); - break; - case 10: - LAUNCH_CUSTOM_ATTENTION_MFMA16(10); - break; - case 11: - LAUNCH_CUSTOM_ATTENTION_MFMA16(11); - break; - case 12: - LAUNCH_CUSTOM_ATTENTION_MFMA16(12); - break; - case 13: - LAUNCH_CUSTOM_ATTENTION_MFMA16(13); - break; - case 14: - LAUNCH_CUSTOM_ATTENTION_MFMA16(14); - break; - case 15: - LAUNCH_CUSTOM_ATTENTION_MFMA16(15); - break; - case 16: - LAUNCH_CUSTOM_ATTENTION_MFMA16(16); - break; - default: - TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); - break; - } - - dim3 reduce_grid(num_heads, num_seqs); - dim3 reduce_block(head_size); - const int warp_size = 32; - const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, warp_size); - // reduction kernel supports upto 16 NPAR_loops * 32 (warp_size) * 256 - // (partition size) = 128K context length - switch (npar_loops) { - case 1: - LAUNCH_CUSTOM_REDUCTION(1); - break; - case 2: - LAUNCH_CUSTOM_REDUCTION(2); - break; - case 3: - LAUNCH_CUSTOM_REDUCTION(3); - break; - case 4: - LAUNCH_CUSTOM_REDUCTION(4); - break; - case 5: - LAUNCH_CUSTOM_REDUCTION(5); - break; - case 6: - LAUNCH_CUSTOM_REDUCTION(6); - break; - case 7: - LAUNCH_CUSTOM_REDUCTION(7); - break; - case 8: - LAUNCH_CUSTOM_REDUCTION(8); - break; - case 9: - LAUNCH_CUSTOM_REDUCTION(9); - break; - case 10: - LAUNCH_CUSTOM_REDUCTION(10); - break; - case 11: - LAUNCH_CUSTOM_REDUCTION(11); - break; - case 12: - LAUNCH_CUSTOM_REDUCTION(12); - break; - case 13: - LAUNCH_CUSTOM_REDUCTION(13); - break; - case 14: - LAUNCH_CUSTOM_REDUCTION(14); - break; - case 15: - LAUNCH_CUSTOM_REDUCTION(15); - break; - case 16: - LAUNCH_CUSTOM_REDUCTION(16); - break; - default: - TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); - break; - } -} - -#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ - PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ - if (!is_navi) { \ - paged_attention_custom_launcher( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ - max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ - } else { \ - paged_attention_custom_launcher_navi( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ - max_seq_len, alibi_slopes, k_scale, v_scale); \ - } - -#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ - OUTT, PSIZE, MFMA_TYPE) \ - if (alibi_slopes) { \ - CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ - true, MFMA_TYPE); \ - } else { \ - CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ - false, MFMA_TYPE); \ - } - -#if defined(__HIPCC__) && defined(__gfx90a__) - #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ - MFMA_TYPE) \ - if (fp8_out_scale) { \ - TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a"); \ - } else { \ - CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ - 256, MFMA_TYPE); \ - } -#else - #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ - MFMA_TYPE) \ - if (fp8_out_scale) { \ - CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ - uint8_t, 256, MFMA_TYPE); \ - } else { \ - CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ - 256, MFMA_TYPE); \ - } -#endif - -#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \ - switch (block_size) { \ - case 16: \ - CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \ - break; \ - case 32: \ - CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \ - switch (head_size) { \ - case 64: \ - CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE); \ - break; \ - case 128: \ - CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported head size: ", head_size); \ - break; \ - } + const std::optional& query_start_loc, int64_t block_size, + int64_t max_seq_len, const std::optional& alibi_slopes, + torch::Tensor& k_scale, torch::Tensor& v_scale, + const std::optional& fp8_out_scale, + const std::string& mfma_type, bool is_navi); bool is_navi_gpu() { static bool is_cached = false; @@ -3646,7 +69,7 @@ void paged_attention( torch::Tensor& query, // [num_seqs, num_heads, head_size] torch::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] torch::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, + int64_t num_kv_heads, double scale, torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] torch::Tensor& seq_lens, // [num_seqs] @@ -3659,48 +82,19 @@ void paged_attention( const std::string& mfma_type) { // clang-format on bool is_navi = is_navi_gpu(); - const int head_size = query.size(2); if (kv_cache_dtype == "auto") { - if (query.dtype() == at::ScalarType::Half) { - CALL_CUSTOM_LAUNCHER_BLK_HEAD( - _Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, MFMAType::F16); - } else if (query.dtype() == at::ScalarType::BFloat16) { - CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16, - vllm::Fp8KVCacheDataType::kAuto, - MFMAType::F16); - } else { - TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); - } + paged_attention_dispatch_auto(out, exp_sums, max_logits, tmp_out, query, + key_cache, value_cache, num_kv_heads, scale, + block_tables, seq_lens, query_start_loc, + block_size, max_seq_len, alibi_slopes, + k_scale, v_scale, fp8_out_scale, is_navi); } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { - if (query.dtype() == at::ScalarType::Half) { - if (mfma_type == "fp8") { - CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t, - vllm::Fp8KVCacheDataType::kFp8E4M3, - MFMAType::Fp8); - } else { - CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t, - vllm::Fp8KVCacheDataType::kFp8E4M3, - MFMAType::F16); - } - } else if (query.dtype() == at::ScalarType::BFloat16) { - if (mfma_type == "fp8") { - CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, - vllm::Fp8KVCacheDataType::kFp8E4M3, - MFMAType::Fp8); - } else { - CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, - vllm::Fp8KVCacheDataType::kFp8E4M3, - MFMAType::F16); - } - } else { - TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); - } + paged_attention_dispatch_fp8(out, exp_sums, max_logits, tmp_out, query, + key_cache, value_cache, num_kv_heads, scale, + block_tables, seq_lens, query_start_loc, + block_size, max_seq_len, alibi_slopes, k_scale, + v_scale, fp8_out_scale, mfma_type, is_navi); } else { TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype); } } - -#undef WARP_SIZE -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/rocm/attention/dispatch_auto.cu b/csrc/rocm/attention/dispatch_auto.cu new file mode 100644 index 000000000000..dc6ac9d70549 --- /dev/null +++ b/csrc/rocm/attention/dispatch_auto.cu @@ -0,0 +1,110 @@ +// Dispatch shard for kv_cache_dtype == "auto" (bf16/fp16 KV cache). +// Split from attention.cu to parallelize compilation across dtype combos. + +#include "launcher.h" + +#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ + PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ + if (!is_navi) { \ + paged_attention_custom_launcher( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ + } else { \ + paged_attention_custom_launcher_navi( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale); \ + } + +#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + OUTT, PSIZE, MFMA_TYPE) \ + if (alibi_slopes) { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + true, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + false, MFMA_TYPE); \ + } + +#if defined(__HIPCC__) && defined(__gfx90a__) + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a"); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#else + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + uint8_t, 256, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#endif + +#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \ + switch (block_size) { \ + case 16: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \ + break; \ + case 32: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ + } + +#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \ + switch (head_size) { \ + case 64: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE); \ + break; \ + case 128: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported head size: ", head_size); \ + break; \ + } + +// clang-format off +void paged_attention_dispatch_auto( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, int64_t num_kv_heads, double scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, + int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + torch::Tensor& k_scale, torch::Tensor& v_scale, + const std::optional& fp8_out_scale, + bool is_navi) { + // clang-format on + const int head_size = query.size(2); + if (query.dtype() == at::ScalarType::Half) { + CALL_CUSTOM_LAUNCHER_BLK_HEAD( + _Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, MFMAType::F16); + } else if (query.dtype() == at::ScalarType::BFloat16) { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16, + vllm::Fp8KVCacheDataType::kAuto, + MFMAType::F16); + } else { + TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); + } +} + +#undef CALL_CUSTOM_LAUNCHER +#undef CALL_CUSTOM_LAUNCHER_ALIBI +#undef CALL_CUSTOM_LAUNCHER_OUT +#undef CALL_CUSTOM_LAUNCHER_BLK +#undef CALL_CUSTOM_LAUNCHER_BLK_HEAD diff --git a/csrc/rocm/attention/dispatch_fp8.cu b/csrc/rocm/attention/dispatch_fp8.cu new file mode 100644 index 000000000000..7ae51d6cf3d2 --- /dev/null +++ b/csrc/rocm/attention/dispatch_fp8.cu @@ -0,0 +1,122 @@ +// Dispatch shard for kv_cache_dtype == "fp8" / "fp8_e4m3" (FP8 KV cache). +// Split from attention.cu to parallelize compilation across dtype combos. + +#include "launcher.h" + +#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ + PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ + if (!is_navi) { \ + paged_attention_custom_launcher( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ + } else { \ + paged_attention_custom_launcher_navi( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale); \ + } + +#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + OUTT, PSIZE, MFMA_TYPE) \ + if (alibi_slopes) { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + true, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + false, MFMA_TYPE); \ + } + +#if defined(__HIPCC__) && defined(__gfx90a__) + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a"); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#else + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + uint8_t, 256, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#endif + +#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \ + switch (block_size) { \ + case 16: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \ + break; \ + case 32: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ + } + +#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \ + switch (head_size) { \ + case 64: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE); \ + break; \ + case 128: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported head size: ", head_size); \ + break; \ + } + +// clang-format off +void paged_attention_dispatch_fp8( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, int64_t num_kv_heads, double scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, + int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + torch::Tensor& k_scale, torch::Tensor& v_scale, + const std::optional& fp8_out_scale, + const std::string& mfma_type, + bool is_navi) { + // clang-format on + const int head_size = query.size(2); + if (query.dtype() == at::ScalarType::Half) { + if (mfma_type == "fp8") { + CALL_CUSTOM_LAUNCHER_BLK_HEAD( + _Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, MFMAType::Fp8); + } else { + CALL_CUSTOM_LAUNCHER_BLK_HEAD( + _Float16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3, MFMAType::F16); + } + } else if (query.dtype() == at::ScalarType::BFloat16) { + if (mfma_type == "fp8") { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::Fp8); + } else { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::F16); + } + } else { + TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); + } +} + +#undef CALL_CUSTOM_LAUNCHER +#undef CALL_CUSTOM_LAUNCHER_ALIBI +#undef CALL_CUSTOM_LAUNCHER_OUT +#undef CALL_CUSTOM_LAUNCHER_BLK +#undef CALL_CUSTOM_LAUNCHER_BLK_HEAD diff --git a/csrc/rocm/attention/kernels.cuh b/csrc/rocm/attention/kernels.cuh new file mode 100644 index 000000000000..bf8137cec738 --- /dev/null +++ b/csrc/rocm/attention/kernels.cuh @@ -0,0 +1,3189 @@ +/* + * Copyright (c) 2024, The vLLM team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include +#include "../../cuda_compat.h" + +#include +#include "../../attention/dtype_fp8.cuh" +#include "../../quantization/w8a8/fp8/amd/quant_utils.cuh" + +// ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent +#if !defined(HIP_FP8_TYPE_OCP) +using __hip_fp8_e4m3 = __hip_fp8_e4m3_fnuz; +using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz; +#endif + +#if defined(__HIPCC__) && \ + (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__GFX9__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__FP8MFMA__ +#endif + +#if defined(NDEBUG) + #undef NDEBUG + #include + #define UNREACHABLE_CODE assert(false); + #define NDEBUG +#else + #define UNREACHABLE_CODE assert(false); +#endif + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) +#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) + +enum class MFMAType { + F16 = 0, + Fp8 = 1, + Fp4 = 2, +}; + +#if defined(__HIP__GFX9__) + + #define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32 + #define GCN_MFMA_INSTR __builtin_amdgcn_mfma_f32_4x4x4f16 + +using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float; +using float16x4 = + __attribute__((__vector_size__(4 * sizeof(_Float16)))) _Float16; +typedef float16x4 _Half4; +using float16x2 = + __attribute__((__vector_size__(2 * sizeof(_Float16)))) _Float16; +typedef float16x2 _Half2; +typedef struct _Half8 { + _Half4 xy[2]; +} _Half8; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; +typedef struct _B16x8 { + _B16x4 xy[2]; +} _B16x8; + +using _B8x8 = uint2; +using _B8x4 = int32_t; // used in builtins +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx4 gcn_mfma4x4x4_instr(const _B16x4& inpA, + const _B16x4& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_4x4x4f16(inpA, inpB, inpC, absz, cbid, + blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(inpA, inpB, inpC, absz, cbid, + blgp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA, + const _B16x4& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x16f16(inpA, inpB, inpC, absz, cbid, + blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(inpA, inpB, inpC, absz, + cbid, blgp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ floatx4 gcn_mfma16x16x32_instr(const long& inpA, + const long& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(inpA, inpB, inpC, absz, + cbid, blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x32_bf8_bf8(inpA, inpB, inpC, absz, + cbid, blgp); + } else { + static_assert(false, "unsupported 8b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x4 from_floatx4(const floatx4& inp) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[2]; + _B16x4 b16x4; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + return u.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + uint32_t u32; + float f32; + } u; + u.f32 = inp[i]; + u.u32 += 0x7fff + ((u.u32 >> 16) & 1); // BF16 RNE with no nan/inf check + ret[i] = uint16_t(u.u32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x4 addx4(const _B16x4& inp1, + const _B16x4& inp2) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + _B16x4 b16x4; + __half2 h2[2]; + } u1, u2, s; + u1.b16x4 = inp1; + u2.b16x4 = inp2; + s.h2[0] = u1.h2[0] + u2.h2[0]; + s.h2[1] = u1.h2[1] + u2.h2[1]; + return s.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + float f32; + uint32_t i32; + } u1, u2, s; + u1.i32 = uint32_t(inp1[i]) << 16; + u2.i32 = uint32_t(inp2[i]) << 16; + s.f32 = u1.f32 + u2.f32; + ret[i] = uint16_t(s.i32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +__device__ __forceinline__ floatx4 to_float_fp8x4(const _B8x4& inp) { + // From MI300+ platforms, we have v_cvt_pk_f32_fp8 instruction + // to convert 2 packed fp8 to 2 packed fp32 values. + // However, in MI200 platforms, we only have v_cvt_f32_fp8 + // to convert fp8 values individually. So we added + // #else case for fewer instructions (# inst=2) in MI300+, + // and fallback to + // #if case for other platforms (# inst=4). + #if defined(__gfx90a__) + float4 f32x4 = vllm::fp8::vec_conversion( + *reinterpret_cast(&inp)); + return *reinterpret_cast(&f32x4); + #else // MI3xx+ optimized builtins + const auto f0 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, false); + const auto f1 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, true); + floatx4 ret; + ret[0] = f0[0]; + ret[1] = f0[1]; + ret[2] = f1[0]; + ret[3] = f1[1]; + return ret; + #endif +} + +template +__device__ __forceinline__ _B16x4 from_floatx4_rtz(const floatx4& inp) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + _Half2 h2[2]; + _B16x4 b16x4; + } u; + u.h2[0] = __builtin_amdgcn_cvt_pkrtz(inp[0], inp[1]); + u.h2[1] = __builtin_amdgcn_cvt_pkrtz(inp[2], inp[3]); + return u.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + uint32_t i32; + float f32; + } u; + u.f32 = inp[i]; + ret[i] = uint16_t(u.i32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) { + union { + _B8x8 b8x8; + _B8x4 b8x4[2]; + } tmp; + tmp.b8x8 = input; + _B16x8 ret; + for (int i = 0; i < 2; i++) { + ret.xy[i] = from_floatx4_rtz(to_float_fp8x4(tmp.b8x4[i])); + } + return ret; +} + +typedef union u64_cvt { + half f16x4[4]; + int16_t b16x4[4]; + _B8x8 b8x8; + _B16x4 b64; + int64_t i64; +} _T8x8; + +__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input, + _T8x8& Mtemp) { + _T8x8 Qtmp8x8; + + for (int i = 0; i < 2; i++) { + floatx4 q_out = {0, 0, 0, 0}; + q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i], + q_out); + Qtmp8x8.b16x4[i * 2] = + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false); + Qtmp8x8.b16x4[i * 2 + 1] = + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false); + } + return Qtmp8x8.b8x8; +} + +__device__ float warpReduceMax(float val) { + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + val = max( + val, __shfl_down(val, offset, WARP_SIZE)); // Using max() for reduction + } + return val; +} + +// grid (num_seqs, num_partitions,num_kv_heads) +// block (256) +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; + const auto warpid = threadIdx.x / WARP_SIZE; + const auto laneid = threadIdx.x % WARP_SIZE; + const int lane4id = laneid % 4; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const auto seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { + return; + } + + const auto partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const auto max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; + + const int partition_start_token_idx = + partition_idx * T_PAR_SIZE; // partition_size; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO4 = DIVIDE_ROUND_UP(GQA_RATIO, 4); + + // shared_logits is used for multiple purposes + __shared__ _B16x4 shared_logits[NWARPS][4][16][4]; + + // for QK mfma16x16, layout is QHead/Tokenx16 across every 16 lanes, 16 Bytes + // HeadElements in each lane, 4x16B HeadElements across 4 rows of warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16; // rows refers to 16 lanes; refer DDP (Data Parallel + // Processing) terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QK_SIZE_RATIO = + sizeof(scalar_t) / + sizeof(cache_t); // 1 for 16bit types, 2 for 8bit types + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 4xQKHE_16B across + // warp + + _B16x8 Qlocal[QKHELOOP] + [QK_SIZE_RATIO]; // note that 16 contiguous elements of Q should + // be fetched per lane for 8 bit cache types : + // QK_SIZE_RATIO changes for this + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each mfma16x16x16 instruction processes 16 tokens + + // can be interpreted as B8x16 for 8 bit types + _B16x8 Klocal[TLOOP][QKHELOOP]; + + const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; + const auto wg_start_kv_head_idx = blockIdx.z; + const auto total_num_heads = gridDim.z * GQA_RATIO; + + // for QK mfma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each mfma takes QH16xT16x16HE across warp + // repeat mfmas across QKHELOOP dimension + // output layout from QKmfma : QH16xT4x4 16 qheads across 16 lanes, 16 tokens + // across 4 rows x 4 tokens per lane + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + #if defined(__HIP__FP8MFMA__) + float q_max = 0; + float q_scale = 1.0; + #endif + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 4 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const int offset1 = + lane16id / + 4; // 16 contiguous chunks of head elems are spread across 4x4lanes + shared_logits[offset1][lane4id][local_qhead_idx][0] = tmp.xy[0]; + shared_logits[offset1][lane4id][local_qhead_idx][1] = tmp.xy[1]; + } else { + for (int i = 0; i < 2; i++) { + const int head_elem = lane16id * 2 + i; // element id in _B16x4 terms + const int offset3 = head_elem % 4; + const int offset2 = (head_elem / 4) % 4; + const int offset1 = head_elem / 4 / 4; + shared_logits[offset1][offset2][local_qhead_idx][offset3] = tmp.xy[i]; + } + } + } + __syncthreads(); + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + for (int i = 0; i < 2; i++) { + Qlocal[qkhe_depth][qkratio].xy[i] = + shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO] + [2 * qkratio + i]; + #if defined(__HIP__FP8MFMA__) + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto && + MFMA_TYPE == MFMAType::Fp8) { + scalar_t* qptr = + reinterpret_cast(&Qlocal[qkhe_depth][qkratio].xy[i]); + for (int k = 0; k < 4; k++) + q_max = fmax(fabs(to_float(qptr[k])), q_max); + } + #endif + } + } + } + + constexpr int KX = + 16 / sizeof(cache_t); // vLLM defines x as 16 Bytes of kv cache elements + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + // fetch K values + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; + } + } + + float alibi_slope; + if constexpr (ALIBI_ENABLED) { + const int alibi_head_idx = wg_start_head_idx + lane16id; + alibi_slope = (lane16id < GQA_RATIO) ? alibi_slopes[alibi_head_idx] : 0.f; + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 64/4 = 16 contiguous vtokens per lane + constexpr int VBLOCKS_PER_LANE = + 1; // assumes block size >=16, each lane can correspond to 1 block only + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = HEAD_SIZE / 16 / NWARPS; + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; + // Safe to use an int32_t here assuming we are working with < 2 billion + // tokens + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + + // v fetches are 16head elems across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int vblock_depth = 0; + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth][vblock_depth]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + } + } + } + + // calculate post qk mfma scale + float scale2 = scale; + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // multiply by k_scale if fp8 kv cache + scale2 *= *k_scale; + #if defined(__HIP__FP8MFMA__) + q_max = warpReduceMax(q_max); + constexpr float FP8_E4M3_SCALE_TARGET = 224.0f; + if constexpr (MFMA_TYPE == MFMAType::Fp8) { + q_scale = q_max > 0 ? FP8_E4M3_SCALE_TARGET / q_max : 1.0f; + scale2 /= q_scale; + } + #endif + } + + floatx4 d_out[TLOOP]; + // qk mfma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + d_out[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + for (int i = 0; i < 2; i++) { + d_out[token_depth] = gcn_mfma16x16x16_instr( + Klocal[token_depth][qkhe_depth].xy[i], + Qlocal[qkhe_depth][qkratio].xy[i], d_out[token_depth]); + } + } + } else { // kv cache dtype fp8 + auto Ktmp = Klocal[token_depth][qkhe_depth]; + _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp); + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + if constexpr (MFMA_TYPE == MFMAType::F16) { + _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio]; + _B16x8 Klocaltmp = convert_b8x8_custom(Ktmp8x8); + for (int i = 0; i < 2; i++) { + d_out[token_depth] = gcn_mfma16x16x16_instr( + Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i], + d_out[token_depth]); + } + } else { + #if defined(__HIP__FP8MFMA__) + _T8x8 Ktmp8x8, Qtmp8x8; + Ktmp8x8.b8x8 = Ktmp8x16.xy[qkratio]; + + for (int n = 0; n < 2; n++) { + scalar_t* qptr = reinterpret_cast( + &Qlocal[qkhe_depth][qkratio].xy[n]); + + Qtmp8x8.b16x4[n * 2] = + vllm::fp8::scaled_vec_conversion( + make_float2(to_float(qptr[0]), + to_float(qptr[1])), + q_scale); + Qtmp8x8.b16x4[n * 2 + 1] = + vllm::fp8::scaled_vec_conversion( + make_float2(to_float(qptr[2]), + to_float(qptr[3])), + q_scale); + } + + d_out[token_depth] = + gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( + Ktmp8x8.i64, Qtmp8x8.i64, d_out[token_depth]); + #else + UNREACHABLE_CODE + #endif + } + } + } + } + d_out[token_depth] *= scale2; + } + + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 4; + + // apply alibi + if constexpr (ALIBI_ENABLED) { + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + const int alibi_offset = local_token_idx - seq_len + 1; + for (int i = 0; i < 4; i++) { + d_out[token_depth][i] += alibi_slope * (alibi_offset + i); + } + } + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 4; i++) { + const float tmp = + (local_token_idx + i < seq_len) ? d_out[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask)); + } + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 4; i++) { + const float tmp = (local_token_idx + i < seq_len) + ? __expf(d_out[token_depth][i] - qk_max) + : 0.0f; + d_out[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { + exp_sum += __shfl_xor(exp_sum, mask); + } + + __syncthreads(); // sync before writing to shared mem + + float* shared_mem = reinterpret_cast(shared_logits); + if (laneid < 16) { + const int qk_max_offset = warpid * 16 + lane16id; + shared_mem[qk_max_offset] = qk_max; + const int exp_sum_offset = NWARPS * 16 + qk_max_offset; + shared_mem[exp_sum_offset] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_mem[w * 16 + lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += + shared_mem[NWARPS * 16 + w * 16 + lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // disable rtz conversion due to its impact on accuracy. + constexpr bool LOGITS_RTZ_CONVERSION = false; + + #if defined(__HIP__FP8MFMA__) + int rowid_8x8 = rowid / 2; + int offset = rowid % 2; + #endif + + // write logits to shared mem + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + d_out[token_depth] *= inv_sum_scale; + if constexpr (MFMA_TYPE != MFMAType::Fp8) { + if constexpr (LOGITS_RTZ_CONVERSION) { + // use rtz conversion for better performance, with negligible impact on + // accuracy + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx4_rtz(d_out[token_depth]); + } else { + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx4(d_out[token_depth]); + } + } else { + #if defined(__HIP__FP8MFMA__) + // cast _B16x4* to _B8x8* + _T8x8& logits_8x8 = *reinterpret_cast<_T8x8*>( + &shared_logits[warpid][token_depth][lane16id][rowid_8x8]); + logits_8x8.b16x4[offset * 2] = __builtin_amdgcn_cvt_pk_fp8_f32( + d_out[token_depth][0], d_out[token_depth][1], 0, false); + logits_8x8.b16x4[offset * 2 + 1] = __builtin_amdgcn_cvt_pk_fp8_f32( + d_out[token_depth][2], d_out[token_depth][3], 0, false); + #else + UNREACHABLE_CODE + #endif + } + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int64_t offset = static_cast(seq_idx) * + static_cast(total_num_heads) * + static_cast(max_num_partitions) + + (static_cast(wg_start_head_idx) + + static_cast(qhead_idx)) * + static_cast(max_num_partitions) + + static_cast(partition_idx); + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + constexpr int ELEMS8_ELEMS4_RATIO = 8 / 4; + constexpr int ELEMS16_ELEMS8_RATIO = 16 / 8; + + _B16x4 outelems[VHELOOP]; + // Softmax V mfma + // v layout: 16he across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx4 tmp_out = {0}; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { + const int offset = rowid * VTLANELOOP * ELEMS8_ELEMS4_RATIO + + vfetch_depth * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems spread + // across 4 rows + tmp_out = gcn_mfma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].xy[i], + shared_logits[vtoken_depth][offset2][lane16id][offset1], + tmp_out); + } + } + // KV cache fp8 + } else { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + _B16x8 Vtmp = Vlocal[vtoken_depth][vhe_depth][vfetch_depth]; + // reinterpret V format as 16 elements of 8bits + _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp); + for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) { + _B8x8 Vtmp8x8 = Vtmp8x16.xy[j]; + if constexpr (MFMA_TYPE == MFMAType::F16) { + _B16x8 Vlocaltmp = convert_b8x8_custom(Vtmp8x8); + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { + const int offset = + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + + j * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems + // spread across 4 rows + tmp_out = gcn_mfma16x16x16_instr( + Vlocaltmp.xy[i], + shared_logits[vtoken_depth][offset2][lane16id][offset1], + tmp_out); + } + } else { + #if defined(__HIP__FP8MFMA__) + for (int i = 0; i < ELEMS8_ELEMS4_RATIO / 2; i++) { + const int offset = + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + + j * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = (offset % ROWS_PER_WARP) / 2; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems + // spread across 4 rows + tmp_out = gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( + reinterpret_cast<_T8x8*>(&Vtmp8x8)->i64, + reinterpret_cast<_T8x8*>( + &shared_logits[vtoken_depth][offset2][lane16id] + [offset1]) + ->i64, + tmp_out); + } + #else + UNREACHABLE_CODE + #endif + } + } + } + } + } + // apply post Softmax V mfma v_scale + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + tmp_out *= *v_scale; + } + outelems[vhe_depth] = from_floatx4(tmp_out); + } + + __syncthreads(); + + // store Softmax-V mfma output to shared mem + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + // lane16 id head dimension; rowid head element dimension + shared_logits[warpid][vhe_depth][lane16id][rowid] = outelems[vhe_depth]; + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO4]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO4; h++) { + const int local_head_idx = 4 * h + rowid; + const int offset1 = (head_elem_idx / 16) % 4; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 4) % 4; + for (int i = 0; i < 2; i++) { + vout[h].xy[i] = + shared_logits[offset1][offset2][local_head_idx][offset3 + i]; + } + } + + const int64_t hsz_maxp_mult = + static_cast(HEAD_SIZE * max_num_partitions); + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO4; h++) { + const int local_head_idx = 4 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int64_t out_head_idx = + static_cast(wg_start_head_idx + local_head_idx); + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +// grid (num_seqs, num_partitions, num_kv_heads) +// block (256 : partition size) +// each WG handles 1 partition per sequence +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; + const auto warpid = threadIdx.x / WARP_SIZE; + const auto laneid = threadIdx.x % WARP_SIZE; + const int lane4id = laneid % 4; + + const auto seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + const auto partition_idx = blockIdx.y; + const auto partition_size = blockDim.x; + const auto max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; + const int partition_start_token_idx = partition_idx * partition_size; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + // every 4 lanes fetch 4 different qheads + // qhloop = num loops over qhead dimension + constexpr int QHLOOP = DIVIDE_ROUND_UP(GQA_RATIO, 4); + constexpr int GQA_RATIO4 = 4 * QHLOOP; + __shared__ float shared_qk_max[NWARPS][GQA_RATIO4 + 1]; + __shared__ float shared_exp_sum[NWARPS][GQA_RATIO4 + 1]; + _B16x8 Qlocal[QHLOOP]; + constexpr int x = 16 / sizeof(scalar_t); + // kheloop = num loops over head_size for 16Bytes of Q/dequantized K elements + constexpr int KHELOOP = HEAD_SIZE / x; + _B16x8 Klocal[KHELOOP]; + _B8x8 Klocalb8[KHELOOP]; + // for SoftMax-V Gemm, V head_size dimension is distributed across warp + // vheloop = num loops to cover v head size dimension + constexpr int VHELOOP = HEAD_SIZE / WARP_SIZE; + // softmax out has warp_size tokens across warp + // vtloop = num loops to cover warp_size(64) tokens with 16Bytes of + // dequantized V elements + constexpr int VTLOOP = WARP_SIZE / 8; + // num vblocks to cover warp_size(64) v elements + constexpr int VBLOCKS = 8 * VTLOOP / BLOCK_SIZE; + int vphysical_blocks[VBLOCKS]; + _B16x8 Vlocal[VHELOOP][VTLOOP]; + _B8x8 Vlocalb8[VHELOOP][VTLOOP]; + floatx4 d_out[QHLOOP]; + float qk_max[QHLOOP]; + + __shared__ _B16x4 vout_shared[QHLOOP][VHELOOP][WARP_SIZE][NWARPS + 1]; + + for (int h = 0; h < QHLOOP; h++) { + d_out[h] = {0}; + qk_max[h] = -FLT_MAX; + } + + const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; + const auto wg_start_kv_head_idx = blockIdx.z; + + const int warp_start_token_idx = + partition_start_token_idx + warpid * WARP_SIZE; + + if (warp_start_token_idx >= seq_len) { // warp out of context + #pragma unroll + for (int h = 0; h < GQA_RATIO4; h++) { + shared_qk_max[warpid][h] = -FLT_MAX; + shared_exp_sum[warpid][h] = 0.0f; + } + } else { // warp within context + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; + // token id within partition + const auto local_token_idx = threadIdx.x; + // token id within sequence + const int global_token_idx = partition_start_token_idx + local_token_idx; + + // fetch block number for k + const int block_idx = (global_token_idx < seq_len) + ? global_token_idx / BLOCK_SIZE + : last_seq_block; + + // fetch k physical block number + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t physical_block_number = + static_cast(block_table[block_idx]); + + // fetch vphysical block numbers up front + const int warp_start_block_idx = warp_start_token_idx / BLOCK_SIZE; + for (int b = 0; b < VBLOCKS; b++) { + const int vblock_idx = warp_start_block_idx + b; + const int vblock_idx_ctx = + (vblock_idx <= last_seq_block) ? vblock_idx : last_seq_block; + vphysical_blocks[b] = block_table[vblock_idx_ctx]; + } + + // fetch q elements + // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elemsc + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + const scalar_t* q_ptr = + q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE; + const _B16x8* q_ptrh8 = reinterpret_cast(q_ptr); + const int qhead_elemh8 = laneid / 4; + + for (int h = 0; h < QHLOOP - 1; h++) { + const int qhead_idx = h * 4 + lane4id; + Qlocal[h] = q_ptrh8[qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; + } + const int final_qhead_idx = 4 * (QHLOOP - 1) + lane4id; + if (final_qhead_idx < GQA_RATIO) { + Qlocal[QHLOOP - 1] = + q_ptrh8[final_qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; + } else { + Qlocal[QHLOOP - 1].xy[0] = {0}; + Qlocal[QHLOOP - 1].xy[1] = {0}; + } + + // fetch k elements + const cache_t* k_ptr = k_cache + physical_block_number * kv_block_stride + + wg_start_kv_head_idx * kv_head_stride; + + // physical_block_offset is already cast in terms of _B16x8 + const int physical_block_offset = local_token_idx % BLOCK_SIZE; + + // each K fetch is for 8 elements of cache_t which are later dequantized to + // scalar_t for fp8 + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const _B16x8* k_ptrh8 = reinterpret_cast(k_ptr); + for (int d = 0; d < KHELOOP; d++) { + Klocal[d] = k_ptrh8[d * BLOCK_SIZE + physical_block_offset]; + } + } else { + // vllm defines X as 16 Bytes of elements of cache_t + constexpr int X = 16 / sizeof(cache_t); + const cache_t* k_ptr2 = k_ptr + physical_block_offset * X; + for (int d = 0; d < KHELOOP; d++) { + const int head_elem = d * 8; + const int offset1 = head_elem / X; + const int offset2 = head_elem % X; + const cache_t* k_ptr3 = k_ptr2 + offset1 * BLOCK_SIZE * X + offset2; + Klocalb8[d] = *reinterpret_cast(k_ptr3); + } + } + + // optional alibi fetch + float alibi_slope[QHLOOP]; + if constexpr (ALIBI_ENABLED) { + for (int h = 0; h < QHLOOP; h++) { + const int qhead_idx = h * 4 + lane4id; + alibi_slope[h] = (qhead_idx < GQA_RATIO) + ? alibi_slopes[wg_start_head_idx + qhead_idx] + : 0.f; + } + } + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; + // fetch vcache in kv cache auto case + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const _B16x8* v_ptrh8 = reinterpret_cast(v_ptr); + // iterate over each v block + for (int b = 0; b < VBLOCKS; b++) { + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t vphysical_block_number = + static_cast(vphysical_blocks[b]); + const _B16x8* v_ptrh8b = + v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; + // iterate over each head elem (within head_size) + for (int h = 0; h < VHELOOP; h++) { + const int head_size_elem = h * WARP_SIZE + laneid; + const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + // iterate over all velems within block + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } + } + } + } // if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) + // fetch vcache in fp8 case + else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) + const _B8x8* v_ptrh8 = reinterpret_cast(v_ptr); + // iterate over each v block + for (int b = 0; b < VBLOCKS; b++) { + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t vphysical_block_number = + static_cast(vphysical_blocks[b]); + const _B8x8* v_ptrh8b = + v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; + // iterate over each head elem (within head_size) + for (int h = 0; h < VHELOOP; h++) { + const int head_size_elem = h * WARP_SIZE + laneid; + const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + // iterate over all velems within block + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } + } + } + } + + #define QK_mfma(x) \ + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ + Klocal[x] = convert_b8x8_custom(Klocalb8[x]); \ + } \ + for (int h = 0; h < QHLOOP; h++) { \ + d_out[h] = gcn_mfma4x4x4_instr( \ + Qlocal[h].xy[0], Klocal[x].xy[0], d_out[h]); \ + d_out[h] = gcn_mfma4x4x4_instr( \ + Qlocal[h].xy[1], Klocal[x].xy[1], d_out[h]); \ + } + // QK mfma with Q mfma block broadcast + // Q values across head_size dimension stored across lanes + // K values across head_size dimension are stored depthwise within lane + // Q broadcast with absz, cbid of mfma instruction + QK_mfma(0); + QK_mfma(1); + QK_mfma(2); + QK_mfma(3); + QK_mfma(4); + QK_mfma(5); + QK_mfma(6); + QK_mfma(7); + // below only needed for head size 128 + if constexpr (KHELOOP > 8) { + QK_mfma(8); + QK_mfma(9); + QK_mfma(10); + QK_mfma(11); + QK_mfma(12); + QK_mfma(13); + QK_mfma(14); + QK_mfma(15); + } + #undef QK_mfma + + float scale2 = scale; + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // post mfma scaling for fp8 + scale2 *= *k_scale; + } + + for (int h = 0; h < QHLOOP; h++) { + d_out[h] *= scale2; + } + + // transpose d_out so that 4 token ids are in each lane, and 4 heads are + // across 4 lanes + for (int h = 0; h < QHLOOP; h++) { + floatx4 tmp = {0}; + for (int i = 0; i < 4; i++) { + const float B = (lane4id == i) ? 1.0f : 0.0f; + tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(d_out[h][i], B, tmp, 0, 0, 0); + } + d_out[h] = tmp; + } + + const int lane4_token_idx = 4 * (global_token_idx >> 2); + + if constexpr (ALIBI_ENABLED) { + const int alibi_offset = lane4_token_idx - seq_len + 1; + for (int h = 0; h < QHLOOP; h++) { + for (int i = 0; i < 4; i++) { + d_out[h][i] += alibi_slope[h] * (alibi_offset + i); + } + } + } + + const int bpermute_mask = 4 * (16 * ((laneid >> 2) % 4) + lane4id); + + for (int h = 0; h < QHLOOP; h++) { + qk_max[h] = -FLT_MAX; + for (int i = 0; i < 4; i++) { + qk_max[h] = (lane4_token_idx + i < seq_len) + ? fmaxf(qk_max[h], d_out[h][i]) + : qk_max[h]; + } + + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { + // qk_max[h] = fmaxf(qk_max[h], __shfl_xor(qk_max[h], mask)); + // } + // faster version of above code with dpp + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + + auto tmp = __builtin_amdgcn_ds_bpermute( + bpermute_mask, *reinterpret_cast(&qk_max[h])); + qk_max[h] = *reinterpret_cast(&tmp); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + } + + float exp_sum[QHLOOP]; + for (int h = 0; h < QHLOOP; h++) { + exp_sum[h] = 0.0f; + for (int i = 0; i < 4; i++) { + d_out[h][i] = (lane4_token_idx + i < seq_len) + ? __expf(d_out[h][i] - qk_max[h]) + : 0.0f; + exp_sum[h] += d_out[h][i]; + } + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { + // exp_sum[h] += __shfl_xor(exp_sum[h], mask); + // } + // faster version of above code with dpp + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + + auto tmp = __builtin_amdgcn_ds_bpermute( + bpermute_mask, *reinterpret_cast(&exp_sum[h])); + exp_sum[h] = *reinterpret_cast(&tmp); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + } + + if (laneid < 4) { + for (int h = 0; h < QHLOOP; h++) { + const int head_idx = 4 * h + lane4id; + shared_qk_max[warpid][head_idx] = qk_max[h]; + shared_exp_sum[warpid][head_idx] = exp_sum[h]; + } + } + } // warp within context + + __syncthreads(); + + const auto num_heads = gridDim.z * GQA_RATIO; + float* max_logits_ptr = + max_logits + seq_idx * num_heads * max_num_partitions + partition_idx; + float* exp_sums_ptr = + exp_sums + seq_idx * num_heads * max_num_partitions + partition_idx; + // calculate qk_max and exp_sums for partition + for (int h = 0; h < QHLOOP; h++) { + float global_qk_max = -FLT_MAX; + float warp_qk_max[NWARPS]; + const int head_idx = 4 * h + lane4id; + for (int w = 0; w < NWARPS; w++) { + warp_qk_max[w] = shared_qk_max[w][head_idx]; + global_qk_max = fmaxf(global_qk_max, warp_qk_max[w]); + } + float global_exp_sum = 0.0f; + for (int w = 0; w < NWARPS; w++) { + global_exp_sum += + shared_exp_sum[w][head_idx] * __expf(warp_qk_max[w] - global_qk_max); + } + if (head_idx < GQA_RATIO) { + max_logits_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = + global_qk_max; + exp_sums_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = + global_exp_sum; + } + const float global_inv_sum_scale = __fdividef(1.f, global_exp_sum + 1e-6f) * + __expf(qk_max[h] - global_qk_max); + d_out[h] *= global_inv_sum_scale; + } + constexpr bool LOGITS_RTZ_CONVERSION = false; + // logits[h] -> every 4 lanes hold 4 heads, each lane holds 4 tokens, there + // are 4x16 tokens across warp + _B16x4 logits[QHLOOP]; + for (int h = 0; h < QHLOOP; h++) { + if constexpr (LOGITS_RTZ_CONVERSION) { + // use rtz for faster performance with no perceivable accuracy loss + logits[h] = from_floatx4_rtz(d_out[h]); + } else { + logits[h] = from_floatx4(d_out[h]); + } + } + + if (warp_start_token_idx >= seq_len) { // warp out of context + for (int qh = 0; qh < QHLOOP; qh++) { + for (int vh = 0; vh < VHELOOP; vh++) { + vout_shared[qh][vh][laneid][warpid] = {0}; + } + } + } else { // warp in context + #define SV_mfma(x) \ + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ + Vlocal[vh][x] = convert_b8x8_custom(Vlocalb8[vh][x]); \ + } \ + for (int qh = 0; qh < QHLOOP; qh++) { \ + acc[qh] = gcn_mfma4x4x4_instr( \ + logits[qh], Vlocal[vh][x].xy[0], acc[qh]); \ + acc[qh] = gcn_mfma4x4x4_instr( \ + logits[qh], Vlocal[vh][x].xy[1], acc[qh]); \ + } + + for (int vh = 0; vh < VHELOOP; vh++) { + floatx4 acc[QHLOOP]; + for (int qh = 0; qh < QHLOOP; qh++) { + acc[qh] = {0}; + } + // SoftMax-V calculation + // logits -> token dimension is distributed across lanes + // Vlocal -> token dimension is depthwise within lane + // uses mfma instruction block broadcast for logits + SV_mfma(0); + SV_mfma(1); + SV_mfma(2); + SV_mfma(3); + SV_mfma(4); + SV_mfma(5); + SV_mfma(6); + SV_mfma(7); + + for (int qh = 0; qh < QHLOOP; qh++) { + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // post mfma v scale for fp8 + acc[qh] *= *v_scale; + } + vout_shared[qh][vh][laneid][warpid] = from_floatx4(acc[qh]); + } + } + + #undef SV_mfma + } // warp in context + + __syncthreads(); + + // final write to tmp_out after vout accumulation + if (warpid == 0) { + _B16x4 vout[QHLOOP][VHELOOP]; + // iterate across heads + for (int qh = 0; qh < QHLOOP; qh++) { + // iterate over each v head elem (within head_size) + for (int vh = 0; vh < VHELOOP; vh++) { + vout[qh][vh] = {0}; + for (int w = 0; w < NWARPS; w++) { + vout[qh][vh] = + addx4(vout[qh][vh], vout_shared[qh][vh][laneid][w]); + } + } + } + + scalar_t* out_ptr = out + + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + partition_idx * HEAD_SIZE; + const int out_num_partitions = max_num_partitions; + bit16_t* out_ptr_b16 = reinterpret_cast(out_ptr); + for (int qh = 0; qh < QHLOOP; qh++) { + for (int vh = 0; vh < VHELOOP; vh++) { + const int head_size_elem = vh * WARP_SIZE + laneid; + for (int i = 0; i < 4; i++) { + const int head_idx = 4 * qh + i; + if (head_idx < GQA_RATIO) { + out_ptr_b16[(wg_start_head_idx + head_idx) * out_num_partitions * + HEAD_SIZE + + head_size_elem] = vout[qh][vh][i]; + } + } + } + } + } // warpid == 0 +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const auto warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 64; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + const float out_scale = + (fp8_out_scale_ptr != nullptr) ? 1.0f / (*fp8_out_scale_ptr) : 1.0f; + acc *= inv_global_exp_sum; + acc *= out_scale; + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + if constexpr (std::is_same::value) { + out_ptr[threadIdx.x] = + __hip_cvt_float_to_fp8(acc, vllm::fp8::fp8_type::__default_saturation, + vllm::fp8::fp8_type::__default_interpret); + } else { + out_ptr[threadIdx.x] = from_float(acc); + } +} + +#elif defined(__GFX11__) + +using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; + +using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; +union b16x8_u { + bit16x8 u16x8; + _B16x4 xy[2]; +}; +typedef b16x8_u _B16x8; + +using bit16x16 = + __attribute__((__vector_size__(16 * sizeof(uint16_t)))) uint16_t; +union b16x16_u { + bit16x16 u16x16; + _B16x8 xy[2]; +}; +typedef b16x16_u _B16x16; + +using _B8x8 = uint2; +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x16& inpA, + const bit16x16& inpB, + const floatx8& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(inpA, inpB, inpC); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(inpA, inpB, inpC); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[4]; + _B16x8 b16x8; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); + u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); + return u.b16x8; + } else if constexpr (std::is_same::value) { + union b2cvt { + __hip_bfloat162 b2[4]; + _B16x8 b16x8; + } u; + + u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); + u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); + u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); + u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); + + return u.b16x8; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 + const int warpid = threadIdx.x / WARP_SIZE; + const int laneid = threadIdx.x % WARP_SIZE; + const int lane2id = laneid % 2; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const int seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { + return; + } + + const int partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const int max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; // length of a seq + + const int partition_start_token_idx = partition_idx * T_PAR_SIZE; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); + + __shared__ float shared_qk_max[NWARPS][16 + 1]; + __shared__ float shared_exp_sum[NWARPS][16 + 1]; + // shared_logits is used for multiple purposes + __shared__ _B16x16 shared_logits[NWARPS][2][16][2]; + + // for QK wmma16x16, layout is QHead/Tokenx16 across every 16 lanes, + // 32 Bytes HeadElements in each lane, 2x16B HeadElements across a row of warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16 / 2; // rows refers to 16 lanes; refer dpp terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across + // warp + + _B16x16 Qlocal[QKHELOOP / 2]; // note that 16 contiguous elements of Q should + // be fetched per lane for 16 bit cache types + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each wmma16x16x16 instruction processes 16 tokens + + _B16x16 Klocal[TLOOP] + [QKHELOOP / 2]; // can be interpreted as B8x16 for 8 bit types + + const int wg_start_head_idx = blockIdx.z * GQA_RATIO; + const int wg_start_kv_head_idx = blockIdx.z; + const int total_num_heads = gridDim.z * GQA_RATIO; + + // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each wmma takes QH16xT16x16HE across warp + // repeat wmma across QKHELOOP dimension + // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens + // across 2 rows x 8 tokens per lane + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + + if (GQA_RATIO == 1) { + const int local_qhead_idx = lane16id % GQA_RATIO; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + if (lane16id < GQA_RATIO) { + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH * 2; + const _B16x16* q_fetch_ptr_32B = + reinterpret_cast(q_fetch_ptr); + Qlocal[qkhe_depth] = *q_fetch_ptr_32B; + } + } + } else { + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 2 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + + const int offset1 = + lane16id / + 2; // 16 contiguous chunks of head elems are spread across 8x2lanes + shared_logits[offset1][lane2id][local_qhead_idx][0].xy[0] = tmp; + } + + __syncthreads(); + + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + Qlocal[qkhe_depth].xy[0] = + shared_logits[qkhe_depth][0][lane16id % GQA_RATIO][0].xy[0]; + Qlocal[qkhe_depth].xy[1] = + shared_logits[qkhe_depth][1][lane16id % GQA_RATIO][0].xy[0]; + } + } + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + constexpr int KX = 16 / sizeof(cache_t); + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = 0; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth / 2].xy[qkhe_depth % 2] = *k_fetch_ptr_16B; + } + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 32/1 = 32 vtokens per lane + constexpr int VBLOCKS_PER_LANE = 2; // assumes block size >=16 + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = DIVIDE_ROUND_UP( + (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each + // wmma instr works on 16 head elements + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + vblock_depth * BLOCK_SIZE; + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x16 Vlocal[VTLOOP][VHELOOP] + [VTLANELOOP / 2]; // this can be interpreted as B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; + // v fetches are 16head elems across lanes x (16x2) tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth] + [vfetch_depth / VBLOCKS_PER_LANE]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + + (vfetch_depth % VBLOCKS_PER_LANE) * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2].xy[vfetch_depth % 2] = + *v_fetch_ptr_16B; + } + } + } + + floatx8 dout[TLOOP]; + // qk wmma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + dout[token_depth] = gcn_wmma16x16x16_instr( + Klocal[token_depth][qkhe_depth].u16x16, Qlocal[qkhe_depth].u16x16, + dout[token_depth]); + } + dout[token_depth] *= scale; + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid; + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = + (local_token_idx + 2 * i < seq_len) ? dout[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = (local_token_idx + 2 * i < seq_len) + ? __expf(dout[token_depth][i] - qk_max) + : 0.0f; + dout[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + exp_sum += __shfl_xor(exp_sum, 16); + + __syncthreads(); + + if (laneid < 16) { + shared_qk_max[warpid][lane16id] = qk_max; + shared_exp_sum[warpid][lane16id] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + #pragma unroll + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // write logits to shared mem + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] *= inv_sum_scale; + shared_logits[warpid][token_depth][lane16id][0].xy[rowid] = + from_floatx8(dout[token_depth]); + } + __syncthreads(); + + _B16x8 swp_buf[TLOOP][2]; + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + swp_buf[token_depth][0] = + shared_logits[warpid][token_depth][lane16id][0].xy[0]; + swp_buf[token_depth][1] = + shared_logits[warpid][token_depth][lane16id][0].xy[1]; + } + + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + #pragma unroll + for (int i = 0; i < 8; i++) { + shared_logits[warpid][token_depth][lane16id][0].xy[rowid].u16x8[i] = + swp_buf[token_depth][i % 2].u16x8[4 * rowid + (i / 2)]; + } + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int offset = seq_idx * total_num_heads * max_num_partitions + + (wg_start_head_idx + qhead_idx) * max_num_partitions + + partition_idx; + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + _B16x8 outelems[VHELOOP]; + // Softmax V wmma + // v layout: 16he across lanes x (16x2) tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx8 tmp_out = {0}; + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP / 2; + vfetch_depth++) { + const int offset = vfetch_depth; + // if output format is 16 qheads across 16 lanes, 16 head elems spread + // across rows + tmp_out = gcn_wmma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x16, + shared_logits[vtoken_depth][offset][lane16id][0].u16x16, tmp_out); + } + } + outelems[vhe_depth] = from_floatx8(tmp_out); + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid] = + outelems[vhe_depth]; // lane16 id head dimension; rowid head element + // dimension + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + swp_buf[vhe_depth][0] = shared_logits[warpid][vhe_depth][lane16id][0].xy[0]; + swp_buf[vhe_depth][1] = shared_logits[warpid][vhe_depth][lane16id][0].xy[1]; + } + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + #pragma unroll + for (int i = 0; i < 8; i++) { + shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid].u16x8[i] = + swp_buf[vhe_depth][i % 2].u16x8[4 * rowid + (i / 2)]; + } + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO2]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + const int offset1 = (head_elem_idx / 16) % NWARPS; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row + vout[h] = + shared_logits[offset1][offset2][local_head_idx][0].xy[offset3]; + } + + const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int out_head_idx = wg_start_head_idx + local_head_idx; + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const int warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 32; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + acc *= inv_global_exp_sum; + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + out_ptr[threadIdx.x] = from_float(acc); +} + +#elif defined(__GFX12__) + +using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; + +using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; +union b16x8_u { + bit16x8 u16x8; + _B16x4 xy[2]; +}; +typedef b16x8_u _B16x8; + +using _B8x8 = uint2; +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, + const bit16x8& inpB, + const floatx8& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(inpA, inpB, inpC); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float_b16(const bit16_t& inp) { + union tmpcvt { + bit16_t u; + _Float16 f; + __hip_bfloat16 b; + } t16; + t16.u = inp; + if constexpr (std::is_same::value) { + return (float)t16.f; + } else if constexpr (std::is_same::value) { + return __bfloat162float(t16.b); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[4]; + _B16x8 b16x8; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); + u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); + return u.b16x8; + } else if constexpr (std::is_same::value) { + union b2cvt { + __hip_bfloat162 b2[4]; + _B16x8 b16x8; + } u; + + u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); + u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); + u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); + u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); + + return u.b16x8; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 + const int warpid = threadIdx.x / WARP_SIZE; + const int laneid = threadIdx.x % WARP_SIZE; + const int lane2id = laneid % 2; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const int seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + const int partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const int max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; // length of a seq + + const int partition_start_token_idx = partition_idx * T_PAR_SIZE; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); + + __shared__ float shared_qk_max[NWARPS][16 + 1]; + __shared__ float shared_exp_sum[NWARPS][16 + 1]; + // shared_logits is used for multiple purposes + __shared__ _B16x8 shared_logits[NWARPS][2][16][2]; + + // for QK wmma16x16_gfx12, layout is QHead/Tokenx16 across every 16 lanes, + // 16 Bytes HeadElements in each lane, 2x16B HeadElements across 2 rows of + // warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16; // rows refers to 16 lanes; refer dpp terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across + // warp + + _B16x8 Qlocal[QKHELOOP]; // note that 16 contiguous elements of Q should + // be fetched per lane for 16 bit cache types + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each wmma16x16x16 instruction processes 16 tokens + + _B16x8 Klocal[TLOOP] + [QKHELOOP]; // can be interpreted as B8x16 for 8 bit types + + const int wg_start_head_idx = blockIdx.z * GQA_RATIO; + const int wg_start_kv_head_idx = blockIdx.z; + const int total_num_heads = gridDim.z * GQA_RATIO; + + // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each wmma takes QH16xT16x16HE across warp + // repeat wmma across QKHELOOP dimension + // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens + // across 2 rows x 8 tokens per lane + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + + if (GQA_RATIO == 1) { + const int local_qhead_idx = lane16id % GQA_RATIO; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = q + query_start_off * q_stride + + global_qhead_idx * HEAD_SIZE + + rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + if (lane16id < GQA_RATIO) { + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + Qlocal[qkhe_depth] = *q_fetch_ptr_16B; + } + } + } else { + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 2 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + + const int offset1 = + lane16id / + 2; // 16 contiguous chunks of head elems are spread across 8x2lanes + shared_logits[offset1][lane2id][local_qhead_idx][0] = tmp; + } + + __syncthreads(); + + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + Qlocal[qkhe_depth] = + shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO][0]; + } + } + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + constexpr int KX = 16 / sizeof(cache_t); + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; + } + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 32/2 = 16 vtokens per lane + constexpr int VBLOCKS_PER_LANE = 1; // assumes block size >=16 + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = DIVIDE_ROUND_UP( + (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each + // wmma instr works on 16 head elements + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x8 Vlocal[VTLOOP][VHELOOP] + [VTLANELOOP]; // this can be interpreted as B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + + // v fetches are 16head elems across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int vblock_depth = 0; + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth][vblock_depth]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + } + } + } + + floatx8 dout[TLOOP]; + // qk wmma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + dout[token_depth] = gcn_wmma16x16x16_instr( + Klocal[token_depth][qkhe_depth].u16x8, Qlocal[qkhe_depth].u16x8, + dout[token_depth]); + } + dout[token_depth] *= scale; + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 8; + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = + (local_token_idx + i < seq_len) ? dout[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = (local_token_idx + i < seq_len) + ? __expf(dout[token_depth][i] - qk_max) + : 0.0f; + dout[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + exp_sum += __shfl_xor(exp_sum, 16); + + __syncthreads(); + + if (laneid < 16) { + shared_qk_max[warpid][lane16id] = qk_max; + shared_exp_sum[warpid][lane16id] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + #pragma unroll + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // write logits to shared mem + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] *= inv_sum_scale; + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx8(dout[token_depth]); + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int offset = seq_idx * total_num_heads * max_num_partitions + + (wg_start_head_idx + qhead_idx) * max_num_partitions + + partition_idx; + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + _B16x8 outelems[VHELOOP]; + // Softmax V wmma + // v layout: 16he across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx8 tmp_out = {0}; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int offset = rowid * VTLANELOOP + vfetch_depth; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // if output format is 16 qheads across 16 lanes, 16 head elems spread + // across rows + tmp_out = gcn_wmma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x8, + shared_logits[vtoken_depth][offset2][lane16id][offset1].u16x8, + tmp_out); + } + } + outelems[vhe_depth] = from_floatx8(tmp_out); + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + shared_logits[warpid][vhe_depth][lane16id][rowid] = + outelems[vhe_depth]; // lane16 id head dimension; rowid head element + // dimension + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO2]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + const int offset1 = (head_elem_idx / 16) % NWARPS; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row + vout[h] = shared_logits[offset1][offset2][local_head_idx][offset3]; + } + + const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int out_head_idx = wg_start_head_idx + local_head_idx; + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const int warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 32; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + acc *= inv_global_exp_sum; + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + out_ptr[threadIdx.x] = from_float(acc); +} + +#else + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + UNREACHABLE_CODE +} +// clang-format on + +#endif diff --git a/csrc/rocm/attention/launch_macros.h b/csrc/rocm/attention/launch_macros.h new file mode 100644 index 000000000000..030614221a00 --- /dev/null +++ b/csrc/rocm/attention/launch_macros.h @@ -0,0 +1,33 @@ +#pragma once + +// These macros reference kernel names and local variables from the launcher +// scope. They must be #included inside the launcher function body. + +#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \ + paged_attention_ll4mi_QKV_mfma16_kernel \ + <<>>( \ + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ + block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ + max_ctx_blocks, k_scale_ptr, v_scale_ptr); + +#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \ + paged_attention_ll4mi_QKV_mfma4_kernel \ + <<>>( \ + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ + block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ + max_ctx_blocks, k_scale_ptr, v_scale_ptr); + +#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \ + paged_attention_ll4mi_reduce_kernel \ + <<>>( \ + out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \ + query_start_loc_ptr, max_num_partitions, fp8_out_scale_ptr); diff --git a/csrc/rocm/attention/launcher.h b/csrc/rocm/attention/launcher.h new file mode 100644 index 000000000000..574ba252c3b7 --- /dev/null +++ b/csrc/rocm/attention/launcher.h @@ -0,0 +1,320 @@ +#pragma once + +#include "kernels.cuh" +#include "launch_macros.h" + +// GFX9 launcher: uses mfma4 for GQA 1-4, mfma16 for 5-16 +template +void paged_attention_custom_launcher( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, int max_seq_len, + const std::optional& alibi_slopes, torch::Tensor& k_scale, + torch::Tensor& v_scale, const std::optional& fp8_out_scale) { + int num_seqs = block_tables.size(0); + int num_heads = query.size(1); + int head_size = query.size(2); + int max_num_blocks_per_seq = block_tables.size(1); + int q_stride = query.stride(0); + int kv_block_stride = key_cache.stride(0); + int kv_head_stride = key_cache.stride(1); + + const int* query_start_loc_ptr = + query_start_loc + ? reinterpret_cast(query_start_loc.value().data_ptr()) + : nullptr; + + const float* alibi_slopes_ptr = + alibi_slopes + ? reinterpret_cast(alibi_slopes.value().data_ptr()) + : nullptr; + + float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); + float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); + T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); + T* query_ptr = reinterpret_cast(query.data_ptr()); + KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); + KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); + int* block_tables_ptr = block_tables.data_ptr(); + int* seq_lens_ptr = seq_lens.data_ptr(); + const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); + const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); + const auto fp8_out_scale_ptr = + fp8_out_scale + ? static_cast(fp8_out_scale.value().data_ptr()) + : nullptr; + OUTT* out_ptr = reinterpret_cast(out.data_ptr()); + + const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); + + constexpr int PARTITION_SIZE = 256; + const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); + const int gqa_ratio = num_heads / num_kv_heads; + assert(num_heads % num_kv_heads == 0); + assert(head_size == HEAD_SIZE); + + constexpr int NTHR = 256; + dim3 grid(num_seqs, max_num_partitions, num_kv_heads); + dim3 block(NTHR); + const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + switch (gqa_ratio) { + case 1: + LAUNCH_CUSTOM_ATTENTION_MFMA4(1); + break; + case 2: + LAUNCH_CUSTOM_ATTENTION_MFMA4(2); + break; + case 3: + LAUNCH_CUSTOM_ATTENTION_MFMA4(3); + break; + case 4: + LAUNCH_CUSTOM_ATTENTION_MFMA4(4); + break; + case 5: + LAUNCH_CUSTOM_ATTENTION_MFMA16(5); + break; + case 6: + LAUNCH_CUSTOM_ATTENTION_MFMA16(6); + break; + case 7: + LAUNCH_CUSTOM_ATTENTION_MFMA16(7); + break; + case 8: + LAUNCH_CUSTOM_ATTENTION_MFMA16(8); + break; + case 9: + LAUNCH_CUSTOM_ATTENTION_MFMA16(9); + break; + case 10: + LAUNCH_CUSTOM_ATTENTION_MFMA16(10); + break; + case 11: + LAUNCH_CUSTOM_ATTENTION_MFMA16(11); + break; + case 12: + LAUNCH_CUSTOM_ATTENTION_MFMA16(12); + break; + case 13: + LAUNCH_CUSTOM_ATTENTION_MFMA16(13); + break; + case 14: + LAUNCH_CUSTOM_ATTENTION_MFMA16(14); + break; + case 15: + LAUNCH_CUSTOM_ATTENTION_MFMA16(15); + break; + case 16: + LAUNCH_CUSTOM_ATTENTION_MFMA16(16); + break; + default: + TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); + break; + } + + dim3 reduce_grid(num_heads, num_seqs); + dim3 reduce_block(head_size); + const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, WARP_SIZE); + switch (npar_loops) { + case 1: + LAUNCH_CUSTOM_REDUCTION(1); + break; + case 2: + LAUNCH_CUSTOM_REDUCTION(2); + break; + case 3: + LAUNCH_CUSTOM_REDUCTION(3); + break; + case 4: + LAUNCH_CUSTOM_REDUCTION(4); + break; + case 5: + LAUNCH_CUSTOM_REDUCTION(5); + break; + case 6: + LAUNCH_CUSTOM_REDUCTION(6); + break; + case 7: + LAUNCH_CUSTOM_REDUCTION(7); + break; + case 8: + LAUNCH_CUSTOM_REDUCTION(8); + break; + default: + TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); + break; + } +} + +// Navi (GFX11/GFX12) launcher: uses mfma16 for all GQA ratios +template +void paged_attention_custom_launcher_navi( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, int max_seq_len, + const std::optional& alibi_slopes, torch::Tensor& k_scale, + torch::Tensor& v_scale) { + int num_seqs = block_tables.size(0); + int num_heads = query.size(1); + int head_size = query.size(2); + int max_num_blocks_per_seq = block_tables.size(1); + int q_stride = query.stride(0); + int kv_block_stride = key_cache.stride(0); + int kv_head_stride = key_cache.stride(1); + + const int* query_start_loc_ptr = + query_start_loc + ? reinterpret_cast(query_start_loc.value().data_ptr()) + : nullptr; + + const float* alibi_slopes_ptr = nullptr; + + float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); + float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); + T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); + T* query_ptr = reinterpret_cast(query.data_ptr()); + KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); + KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); + int* block_tables_ptr = block_tables.data_ptr(); + int* seq_lens_ptr = seq_lens.data_ptr(); + + const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); + const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); + const auto fp8_out_scale_ptr = nullptr; + OUTT* out_ptr = reinterpret_cast(out.data_ptr()); + + const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); + + constexpr int PARTITION_SIZE = 256; + const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); + const int gqa_ratio = num_heads / num_kv_heads; + assert(num_heads % num_kv_heads == 0); + assert(head_size == HEAD_SIZE); + + constexpr int NTHR = 256; + dim3 grid(num_seqs, max_num_partitions, num_kv_heads); + dim3 block(NTHR); + const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + switch (gqa_ratio) { + case 1: + LAUNCH_CUSTOM_ATTENTION_MFMA16(1); + break; + case 2: + LAUNCH_CUSTOM_ATTENTION_MFMA16(2); + break; + case 3: + LAUNCH_CUSTOM_ATTENTION_MFMA16(3); + break; + case 4: + LAUNCH_CUSTOM_ATTENTION_MFMA16(4); + break; + case 5: + LAUNCH_CUSTOM_ATTENTION_MFMA16(5); + break; + case 6: + LAUNCH_CUSTOM_ATTENTION_MFMA16(6); + break; + case 7: + LAUNCH_CUSTOM_ATTENTION_MFMA16(7); + break; + case 8: + LAUNCH_CUSTOM_ATTENTION_MFMA16(8); + break; + case 9: + LAUNCH_CUSTOM_ATTENTION_MFMA16(9); + break; + case 10: + LAUNCH_CUSTOM_ATTENTION_MFMA16(10); + break; + case 11: + LAUNCH_CUSTOM_ATTENTION_MFMA16(11); + break; + case 12: + LAUNCH_CUSTOM_ATTENTION_MFMA16(12); + break; + case 13: + LAUNCH_CUSTOM_ATTENTION_MFMA16(13); + break; + case 14: + LAUNCH_CUSTOM_ATTENTION_MFMA16(14); + break; + case 15: + LAUNCH_CUSTOM_ATTENTION_MFMA16(15); + break; + case 16: + LAUNCH_CUSTOM_ATTENTION_MFMA16(16); + break; + default: + TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); + break; + } + + dim3 reduce_grid(num_heads, num_seqs); + dim3 reduce_block(head_size); + const int warp_size = 32; + const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, warp_size); + switch (npar_loops) { + case 1: + LAUNCH_CUSTOM_REDUCTION(1); + break; + case 2: + LAUNCH_CUSTOM_REDUCTION(2); + break; + case 3: + LAUNCH_CUSTOM_REDUCTION(3); + break; + case 4: + LAUNCH_CUSTOM_REDUCTION(4); + break; + case 5: + LAUNCH_CUSTOM_REDUCTION(5); + break; + case 6: + LAUNCH_CUSTOM_REDUCTION(6); + break; + case 7: + LAUNCH_CUSTOM_REDUCTION(7); + break; + case 8: + LAUNCH_CUSTOM_REDUCTION(8); + break; + case 9: + LAUNCH_CUSTOM_REDUCTION(9); + break; + case 10: + LAUNCH_CUSTOM_REDUCTION(10); + break; + case 11: + LAUNCH_CUSTOM_REDUCTION(11); + break; + case 12: + LAUNCH_CUSTOM_REDUCTION(12); + break; + case 13: + LAUNCH_CUSTOM_REDUCTION(13); + break; + case 14: + LAUNCH_CUSTOM_REDUCTION(14); + break; + case 15: + LAUNCH_CUSTOM_REDUCTION(15); + break; + case 16: + LAUNCH_CUSTOM_REDUCTION(16); + break; + default: + TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); + break; + } +} diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index ee7f06cbe257..b39862c9123c 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -1,1260 +1,9 @@ -#include -#include -#include - -#include -#include -#include - -#include -#include - -#include "../cuda_compat.h" -#include "dispatch_utils.h" +#define SKINNY_GEMMS_MAIN_TU +#include "skinny_gemms/kernel.cuh" +#include "skinny_gemms/launch.h" #include "quantization/w8a8/fp8/common.cuh" #include "core/batch_invariant.hpp" -// TODO(rasmith): The kernels in this file are susceptible to integer overflow -// issues, do not take strides, and are unable to handle PyTorch tensors that -// return is_contiguous() as False (the tensors may actually be contiguous -// in memory). -// -// However, it may be possible to fix these kernels to handle both issues. - -#if defined(__HIPCC__) && \ - (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) - #define __HIP__GFX9__ -#endif - -// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts -#if defined(__GFX11__) || defined(__GFX12__) - #define __HIP__GFX1X__ -#endif - -#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) - #define __HIP__MI3XX__ -#endif - -#if defined(__gfx950__) - #define LDS_SIZE 160 * 1024 -#else - #define LDS_SIZE 64 * 1024 -#endif - -int get_lds_size() { - static const int result = [] { - const auto* dprops = at::cuda::getCurrentDeviceProperties(); - const std::string device_arch = dprops->gcnArchName; - return device_arch.find("gfx95") == std::string::npos ? 64 * 1024 - : 160 * 1024; - }(); - return result; -} - -bool on_gfx1x() { - static const bool result = [] { - const auto* dprops = at::cuda::getCurrentDeviceProperties(); - const std::string device_arch = dprops->gcnArchName; - return device_arch.find("gfx11") != std::string::npos || - device_arch.find("gfx12") != std::string::npos; - }(); - return result; -} - -bool on_gfx12() { - static const bool result = [] { - const auto* dprops = at::cuda::getCurrentDeviceProperties(); - const std::string device_arch = dprops->gcnArchName; - return device_arch.find("gfx12") != std::string::npos; - }(); - return result; -} - -bool is_gfx11() { - static const bool result = [] { - auto dprops = at::cuda::getCurrentDeviceProperties(); - std::string device_arch = dprops->gcnArchName; - return device_arch.find("gfx11") != std::string::npos; - }(); - return result; -} - -#if defined(NDEBUG) - #undef NDEBUG - #include - #define UNREACHABLE_CODE assert(false); - #define NDEBUG -#else - #define UNREACHABLE_CODE assert(false); -#endif - -template -struct scalar {}; - -template -struct scalar2 {}; - -template -__device__ __forceinline__ float2 __s22float2(T v); - -template -__device__ __forceinline__ T __float2s(float v); - -template -__device__ __forceinline__ T __float22s2_rn(float2 v); - -// Definitions and cvt functions for fp16 -template <> -struct scalar { - using type = half; -}; - -template <> -struct scalar2 { - using type = __half2; -}; - -template <> -__device__ __forceinline__ half __float2s(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ float2 __s22float2(__half2 v) { - return __half22float2(v); -} - -template <> -__device__ __forceinline__ __half2 __float22s2_rn(float2 v) { - return __float22half2_rn(v); -} - -// Definitions and cvt functions for bf16 -template <> -struct scalar { - using type = __hip_bfloat16; -}; - -template <> -struct scalar2 { - using type = __hip_bfloat162; -}; - -template <> -__device__ __forceinline__ __hip_bfloat16 __float2s(float v) { - return __float2bfloat16(v); -} - -template <> -__device__ __forceinline__ float2 __s22float2(__hip_bfloat162 v) { - return __bfloat1622float2(v); -} - -template <> -__device__ __forceinline__ __hip_bfloat162 __float22s2_rn(float2 v) { - return __float22bfloat162_rn(v); -} - -template -__device__ __forceinline__ T loadnt(T* addr) { - return __builtin_nontemporal_load(addr); -} - -__device__ __forceinline__ float4 load_ntmprl(const float4* addr) { - auto addr_alias = reinterpret_cast(addr); - auto dat0 = loadnt(addr_alias); - auto dat1 = loadnt(addr_alias + 1); - auto dat2 = loadnt(addr_alias + 2); - auto dat3 = loadnt(addr_alias + 3); - return make_float4(dat0, dat1, dat2, dat3); -} - -// TBlock fetches entire rows of A, and entire col of B (K dimension); assume -// N=1 for time being grid is M/A_NUM_ROWS blocks -template -__global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b, - scalar_t* out_c, const int K) { - using scalar2_t = typename scalar2::type; - auto af4 = reinterpret_cast(in_a); - auto bf4 = reinterpret_cast(in_b); - auto c = reinterpret_cast(out_c); - __shared__ float red_smem[NUM_A_ROWS_PER_BLOCK][WARP_SIZE]; - const int row_addr = blockIdx.x * NUM_A_ROWS_PER_BLOCK * K / 8; - const int threadid = threadIdx.x; - const int warp = threadIdx.x / WARP_SIZE; - const int lane = threadIdx.x % WARP_SIZE; - const int num_warps = blockDim.x / WARP_SIZE; - const int qwarpid = threadid / 16; - const int qthreadid = threadid % 16; - float4 rowA_elem4[NUM_A_ROWS_PER_BLOCK]; - scalar2_t colB_elem4x, colB_elem4y, colB_elem4z, colB_elem4w; - float acc[NUM_A_ROWS_PER_BLOCK]; - scalar2_t acch2; - scalar2_t oval; - - // Each thread processes 8 elements per iteration. With NUM_THREADS threads, - // each iteration covers NUM_THREADS * 8 elements of K. Loop over chunks - // to handle K values larger than NUM_THREADS * 8. - const int elems_per_iter = blockDim.x; // threads, each handling 8 elements - const int K_div8 = K / 8; - -#pragma unroll - for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { - acc[i] = 0.f; - } - - for (int base = 0; base < K_div8; base += elems_per_iter) { - int idx = base + threadid; - if (idx < K_div8) { -#pragma unroll - for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { - rowA_elem4[i] = load_ntmprl(&af4[row_addr + idx + K_div8 * i]); - } - colB_elem4x = bf4[idx * 4 + 0]; - colB_elem4y = bf4[idx * 4 + 1]; - colB_elem4z = bf4[idx * 4 + 2]; - colB_elem4w = bf4[idx * 4 + 3]; - - scalar2_t Af2; - auto Ah2ptr = reinterpret_cast(&rowA_elem4); - scalar2_t* ah2lptr; - -#pragma unroll - for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { - ah2lptr = Ah2ptr + i * 4; - Af2 = *(ah2lptr); - acch2 = __hmul2(Af2, colB_elem4x); - Af2 = *(ah2lptr + 1); - acch2 = __hfma2(Af2, colB_elem4y, acch2); - Af2 = *(ah2lptr + 2); - acch2 = __hfma2(Af2, colB_elem4z, acch2); - Af2 = *(ah2lptr + 3); - acch2 = __hfma2(Af2, colB_elem4w, acch2); - float2 S = __s22float2(acch2); - acc[i] += S.x + S.y; - } - } - } - -// all reduce across warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { -#pragma unroll - for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { - acc[i] += __shfl_xor(acc[i], mask); - } - } - - // Warp leaders store the data to shared memory. - if (lane < NUM_A_ROWS_PER_BLOCK) { - red_smem[lane][warp] = acc[lane]; - } - - // Make sure the data is in shared memory. - __syncthreads(); - - if (qwarpid < NUM_A_ROWS_PER_BLOCK) { - acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f; -#pragma unroll - for (int mask = 16 / 2; mask >= 1; mask /= 2) { - acc[qwarpid] += __shfl_xor(acc[qwarpid], mask); - } - float oval2 = __shfl_xor(acc[qwarpid], 16); - - if (lane % 32 == 0) { - oval = __float22s2_rn(make_float2(acc[qwarpid], oval2)); - c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval; - } - } -} - -torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, - const int64_t rows_per_block) { - auto M = in_a.size(0); - auto K = in_a.size(1); - auto N = in_b.size(0); - - TORCH_CHECK(N == 1, "Row number of activation tensor must be 1."); - TORCH_CHECK(in_a.dtype() == in_b.dtype()); - TORCH_CHECK(in_b.dtype() == torch::kFloat16 || - in_b.dtype() == torch::kBFloat16); - - auto out_c = torch::empty( - {N, M}, torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device())); - - // NUM_THREADS must be a multiple of WARP_SIZE (warp shuffle operations). - // Cap at 512 (16 warps) because the cross-warp reduction uses 16 threads - // per row. The kernel loops over K in chunks when K/8 > NUM_THREADS. - int NUM_THREADS = - max(rows_per_block * 16, - K * 2 / 16 % WARP_SIZE == 0 - ? K * 2 / 16 - : K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE)); - NUM_THREADS = min(NUM_THREADS, 512); - - int NUM_BLOCKS = M / rows_per_block; - - const at::cuda::OptionalCUDAGuard device_guard(device_of(in_b)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - // call the kernel function... - AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "LLGemm1", [&] { - auto a_ptr = in_a.data_ptr(); - auto b_ptr = in_b.data_ptr(); - auto c_ptr = out_c.data_ptr(); - if (rows_per_block == 2) { - LLGemm1_kernel - <<>>(a_ptr, b_ptr, c_ptr, K); - } else if (rows_per_block == 4) { - LLGemm1_kernel - <<>>(a_ptr, b_ptr, c_ptr, K); - } else if (rows_per_block == 8) { - LLGemm1_kernel - <<>>(a_ptr, b_ptr, c_ptr, K); - } else if (rows_per_block == 16) { - LLGemm1_kernel - <<>>(a_ptr, b_ptr, c_ptr, K); - } else { - NUM_BLOCKS = M / 4; - LLGemm1_kernel - <<>>(a_ptr, b_ptr, c_ptr, K); - } - }); - - return out_c; -} - -#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__) - #define DOT2C(V0, V2, V3) \ - if constexpr (std::is_same_v) { \ - asm("v_dot2c_f32_f16 %0, %2, %3" \ - : "=v"(V0) \ - : "0"(V0), "v"(V2), "v"(V3)); \ - } else if constexpr (std::is_same_v) { \ - float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ - __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ - V0 += (s.x + s.y); \ - } -#elif defined(__HIP__GFX1X__) - // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12) - #define DOT2C(V0, V2, V3) \ - if constexpr (std::is_same_v) { \ - asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ - } else if constexpr (std::is_same_v) { \ - float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ - __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ - V0 += (s.x + s.y); \ - } -#endif - -// To avoid LLVM silently upcasting to double -__device__ inline unsigned int min__(uint32_t a, uint32_t b) { - return min(a, b); -} - -#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) -// META3-2: Variant of the wvSplitK_hf_sml LDS-load loop that fuses the -// silu_and_mul preamble. The source activation tensor A has 2*K columns -// per row, packed as [gate(K) | up(K)]; the LDS staging buffer ends up -// with silu(gate) * up. N=1 only (decode batch=1 path). Mirrors the -// int4 EXPERIMENT-g helper `load_act_into_lds_silu_mul` for the bf16 / -// fp16 wvSplitK template. -// -// Match the unfused silu_and_mul semantics exactly: silu in fp32, cast -// back to scalar_t, then the scalar_t multiply by `up`. Doing the -// multiply in fp32 changed downstream GEMM rounding subtly enough to -// regress generated text on the int4 path; same caution applies here. -template -__device__ __forceinline__ void load_act_into_lds_silu_mul_bf16( - scalar_t* s, const scalar_t* __restrict__ A, const int K, - const int max_lds_len) { - using scalar8 = - __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; - union bigType { - scalar_t h[A_CHUNK]; - float f[A_CHUNK / 2]; - scalar8 h8; - }; - const int limit = min__(K, max_lds_len); - for (uint32_t k = 0; k < (uint32_t)limit; k += THRDS * WvPrGrp * A_CHUNK) { - uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); - if (k_in >= (uint32_t)limit) break; - bigType gate = *((const bigType*)(&A[k_in])); - bigType up = *((const bigType*)(&A[k_in + K])); - bigType out; - #pragma unroll - for (int i = 0; i < A_CHUNK; ++i) { - float g; - if constexpr (std::is_same_v) { - g = __bfloat162float(gate.h[i]); - } else { - g = __half2float(gate.h[i]); - } - scalar_t silu_g = __float2s(g / (1.0f + expf(-g))); - out.h[i] = silu_g * up.h[i]; - } - *((bigType*)(&s[k_in])) = out; - } - __syncthreads(); -} - -// This version targets cases where A[] fits LDS capacity -template -__global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, const int M, - const int Bx, const int By, const scalar_t* B, - const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount, - const scalar_t* __restrict__ GATE = nullptr) { - static_assert(!FUSED_SILU_MUL || N == 1, - "FUSED_SILU_MUL is only supported with N=1"); - static_assert(!FUSED_GATE_MUL || N == 1, - "FUSED_GATE_MUL is only supported with N=1"); - constexpr int max_lds_len = LDS_SIZE / 2; - #if defined(__HIP__MI3XX__) - constexpr bool use_mfma = (std::is_same_v); - #else - constexpr bool use_mfma = false; - #endif - using scalar8 = - __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; - using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; - union bigType { - scalar_t h[A_CHUNK]; - float f[A_CHUNK / 2]; - float2 f2[A_CHUNK / 4]; - double d[A_CHUNK / 4]; - half4 h4[A_CHUNK / 4]; - scalar8 h8; - }; - - //---------------------------------------------------- - // Reserving 64/160 KB of LDS to have 1 WG / CU - // Goal is to bring the activation matrix A to the LDS - // and use it across the lifetime of the work group - // TODO: When activation matrix is larger than 64 KB - // then this is not going to work! - //---------------------------------------------------- - __shared__ scalar_t s[max_lds_len]; - - //---------------------------------------------------- - // Fetch the activation matrix to LDS - // Loop iteration: - // - Each thread (lane) is fetching 8 elements (A_Chunk) - // - Each wave will fetch 64*8=> 512 elements - // - Each WG will fetch 512 * 16 => 8K elements - // - Then the WG will move to another 8 K elements - // TODO: Logic below will only work when K is multiple of 8 - //---------------------------------------------------- - if constexpr (FUSED_SILU_MUL) { - // META3-2: A is laid out [gate(K) | up(K)] per row; this writes - // silu(gate)*up into LDS so the kernel sees the post-activation K - // elements directly. N=1 (asserted above). - load_act_into_lds_silu_mul_bf16( - s, A, K, max_lds_len); - } else { - for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; - k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { - #if defined(__gfx950__) - __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); - #else - *((bigType*)(&s[k])) = *((bigType*)(&A[k])); - #endif - } - __syncthreads(); - } - - if (threadIdx.y >= _WvPrGrp) return; - - uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; - - //---------------------------------------------------- - // Each wave works on a single column of weight matrix. - // There are 16 waves per WG, and hence, each WG is - // working on 16 columns of weight matrix. Moreover, - // we tile in column direction by YTILE, so when YTILE=1 - // the above math is right, however, when YTILE=2 then - // each wave will be working on 2 columns and WG will - // be working on 32 columns. - // - // Top level loop that makes WGs persistent! - // - WGs iterates across columns of weight matrix - // - Each wave within WG works on a given column(s) - // - After completing first set of columns, WGs start - // working on the next set of available columns - //---------------------------------------------------- - while (m < M) { - //---------------------------------------------------- - // 'sum' accumulates the matrix A x B computation - // split across 64 lanes. - // - // YTILE represents how many column of weight matrix - // are being worked on by each wave. - //---------------------------------------------------- - float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; - - for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; - bigType bigB[YTILE][UNRL]; - // Fetch the weight matrix from memory! - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; - for (int y = 0; y < YTILE; y++) - bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); - } - - // Fetch activation matrix from either just LDS or from both LDS / memory - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - } - } - - // Do the matrix multiplication in interleaved manner - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); - } - } - } - } - __builtin_amdgcn_sched_barrier(0); - //---------------------------------------------------- - // Final reduction step using shuffle - //---------------------------------------------------- - if constexpr (!use_mfma) { - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, - 1); // row_shr8 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, - 1); // row_shr4 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, - 1); // row_shr2 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, - 1); // row_shr1 - #if defined(__HIP__GFX9__) - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - #else - sum[n][y] += __shfl_xor(sum[n][y], 16); - #endif - } - } - - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (std::is_same_v) { - sum[n][y] += __half2float(biases[n][y]); - } else if constexpr (std::is_same_v) { - sum[n][y] += __bfloat162float(biases[n][y]); - } - scalar_t out_val = __float2s(sum[n][y]); - if constexpr (FUSED_GATE_MUL) { - // META3-2 Phase 2: per-token scalar mul (mirrors the - // unfused `F.sigmoid(expert_gate(x)) * out` in scalar_t). - out_val = out_val * GATE[n]; - } - C[m + y + n * M] = out_val; - } - } - } - } else { - #ifdef __HIP__GFX9__ - #pragma unroll - for (int n = 0; n < N; n++) { - #pragma unroll - for (int y = 0; y < YTILE; y++) { - /*float accm1 = 0; - for (int i=0; i<64; i++) - accm1 += __shfl(sum4[n][y][i%4], i); - sum4[n][y][0] = accm1;*/ - float accm = sum4[n][y][0]; - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, - 1); // row_shl1 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, - 1); // row_shl2 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, - 1); // row_shl3 - accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, - 1); // row_shl4 - accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, - 1); // row_shl8 - accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, - 1); // row_shr15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - - sum4[n][y][0] = accm; - } - } - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - sum4[n][y][0] += __bfloat162float(biases[n][y]); - scalar_t out_val = __float2bfloat16(sum4[n][y][0]); - if constexpr (FUSED_GATE_MUL) { - // META3-2 Phase 2: per-token scalar mul. - out_val = out_val * GATE[n]; - } - C[m + y + n * M] = out_val; - } - } - } - #endif // __HIP__GFX9__ (MFMA path) - } - m += CuCount * _WvPrGrp * YTILE; - } -} -#else -template -__global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, - const int M, const int Bx, const int By, - const scalar_t* B, - const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount, - const scalar_t* __restrict__ GATE = nullptr) { - UNREACHABLE_CODE -} -#endif - -#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) -// This version targets cases where A[] marginally exceeds LDS capacity -template -__global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitK_hf_(const int K, const int Kbp, const int Kap, const int M, - const int Bx, const int By, const scalar_t* B, - const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { - constexpr int max_lds_len = LDS_SIZE / 2; - #if defined(__HIP__MI3XX__) - constexpr bool use_mfma = (std::is_same_v); - #else - constexpr bool use_mfma = false; - #endif - - using scalar8 = - __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; - using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; - union bigType { - scalar_t h[A_CHUNK]; - float f[A_CHUNK / 2]; - float2 f2[A_CHUNK / 4]; - double d[A_CHUNK / 4]; - half4 h4[A_CHUNK / 4]; - scalar8 h8; - }; - - __shared__ scalar_t s[max_lds_len]; - - //---------------------------------------------------- - // Computation of columns that need to be committed to memory! - //---------------------------------------------------- - uint32_t commitColumn[YTILE]; - for (uint32_t i = 0; i < YTILE; i++) { - commitColumn[i] = 1; - } - - uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; - - // Check whether there will be fragmentation! - // This will happen only for the last wave! - if (m < M && (m + YTILE) >= M) { - uint32_t startColumn = M - YTILE; - for (uint32_t i = 0; i < (m - startColumn); i++) { - commitColumn[i] = 0; - } - m = startColumn; - } - - for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; - k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { - #if defined(__gfx950__) - __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); - #else - *((bigType*)(&s[k])) = *((bigType*)(&A[k])); - #endif - } - - __syncthreads(); - - if (threadIdx.y >= _WvPrGrp) return; - - while (m < M) { - float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; - - for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; - bigType bigB[YTILE][UNRL]; - // Fetch the weight matrix from memory! - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; - for (int y = 0; y < YTILE; y++) - bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); - } - - // Fetch activation matrix from either just LDS or from both LDS / memory - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); - } - } - - // Do the matrix multiplication in interleaved manner - for (uint32_t n = 0; n < N; n++) { - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); - } - } - } - } - - //---------------------------------------------------- - // Final reduction step using shuffle - //---------------------------------------------------- - if constexpr (!use_mfma) { - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, - 1); // row_shr8 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, - 1); // row_shr4 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, - 1); // row_shr2 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, - 1); // row_shr1 - #if defined(__HIP__GFX9__) - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - #else - sum[n][y] += __shfl_xor(sum[n][y], 16); - #endif - } - } - - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if (commitColumn[y]) { - if constexpr (std::is_same_v) { - sum[n][y] += __half2float(biases[n][y]); - } else if constexpr (std::is_same_v) { - sum[n][y] += __bfloat162float(biases[n][y]); - } - C[m + y + n * M] = __float2s(sum[n][y]); - } - } - } - } - } else { - #ifdef __HIP__GFX9__ - #pragma unroll - for (int n = 0; n < N; n++) { - #pragma unroll - for (int y = 0; y < YTILE; y++) { - // float accm1 = 0; - // for (int i=0; i<64; i++) - // accm1 += __shfl(sum4[n][y][i%4], i); - float accm = sum4[n][y][0]; - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, - 1); // row_shl1 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, - 1); // row_shl2 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, - 1); // row_shl3 - accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, - 1); // row_shl4 - accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, - 1); // row_shl8 - accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, - 1); // row_shr15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - sum4[n][y][0] = accm; - } - } - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if (commitColumn[y]) { - sum4[n][y][0] += __bfloat162float(biases[n][y]); - C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); - } - } - } - } - #endif // __HIP__GFX9__ (MFMA path) - } - - m += CuCount * _WvPrGrp * YTILE; - - // Check whether there will be fragmentation! - // This will happen only for the last wave! - if (m < M && (m + YTILE) >= M) { - uint32_t startColumn = M - YTILE; - for (uint32_t i = 0; i < (m - startColumn); i++) { - commitColumn[i] = 0; - } - m = startColumn; - } - } -} - -#else -template -__global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap, - const int M, const int Bx, const int By, - const scalar_t* B, const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { - UNREACHABLE_CODE -} -#endif - -#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) -// This version targets big A[] cases, where it is much larger than LDS capacity -template -__global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, const int M, - const int Bx, const int By, const scalar_t* B, - const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { - constexpr int max_lds_len = LDS_SIZE / 2; - #if defined(__HIP__MI3XX__) - constexpr bool use_mfma = (std::is_same_v); - #else - constexpr bool use_mfma = false; - #endif - - using scalar8 = - __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; - using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; - union bigType { - scalar_t h[A_CHUNK]; - float f[A_CHUNK / 2]; - float2 f2[A_CHUNK / 4]; - double d[A_CHUNK / 4]; - half4 h4[A_CHUNK / 4]; - scalar8 h8; - }; - - //---------------------------------------------------- - // Reserving 64/160 KB of LDS to have 1 WG / CU - // Goal is to bring the activation matrix A to the LDS - // and use it across the lifetime of the work group - // TODO: When activation matrix is larger than 64 KB - // then this is not going to work! - //---------------------------------------------------- - __shared__ scalar_t s[max_lds_len]; - - //---------------------------------------------------- - // Computation of columns that need to be committed to memory! - //---------------------------------------------------- - uint32_t commitColumn[YTILE]; - for (uint32_t i = 0; i < YTILE; i++) { - commitColumn[i] = 1; - } - - // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp); - if (threadIdx.y >= _WvPrGrp) return; - - //---------------------------------------------------- - // Indexing function into the column of weight matrix B - // Algorithm does 64 lane k-splitting / wave and uses - // WG ID and Thread ID to find the index. - //---------------------------------------------------- - uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; - - // Check whether there will be fragmentation! - // This will happen only for the last wave! - if (m < M && (m + YTILE) >= M) { - uint32_t startColumn = M - YTILE; - for (uint32_t i = 0; i < (m - startColumn); i++) { - commitColumn[i] = 0; - } - m = startColumn; - } - - //---------------------------------------------------- - // Fetch the activation matrix to LDS - // Loop iteration: - // - Each thread (lane) is fetching 8 elements (A_Chunk) - // - Each wave will fetch 64*8=> 512 elements - // - Each WG will fetch 512 * 16 => 8K elements - // - Then the WG will move to another 8 K elements - // TODO: Logic below will only work when K is multiple of 8 - //---------------------------------------------------- - #define PCML - #ifndef PCML - for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; - k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { - #if defined(__gfx950__) - __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); - #else - *((bigType*)(&s[k])) = *((bigType*)(&A[k])); - #endif - } - __syncthreads(); - #endif - - #define TUC (THRDS * UNRL * A_CHUNK) - uint32_t kBase = 0; - // find biggest k size that fits in LDS - uint32_t kFit = (max_lds_len) / N; - // kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple - // of TUC - kFit = (kFit % TUC == 0) - ? kFit - : (kFit - kFit % TUC); // round up to multiple of TUC - // if (kFit == 0) kFit = TUC; - kFit = min__(kFit, Kap); - - //---------------------------------------------------- - // Each wave works on a single column of weight matrix. - // There are 16 waves per WG, and hence, each WG is - // working on 16 columns of weight matrix. Moreover, - // we tile in column direction by YTILE, so when YTILE=1 - // the above math is right, however, when YTILE=2 then - // each wave will be working on 2 columns and WG will - // be working on 32 columns. - // - // Top level loop that makes WGs persistent! - // - WGs iterates across columns of weight matrix - // - Each wave within WG works on a given column(s) - // - After completing first set of columns, WGs start - // working on the next set of available columns - //---------------------------------------------------- - #ifdef PCML - int YW = (YTILE * _WvPrGrp); - uint32_t Mrndp = (M % YW == 0) ? M : (M - M % YW + YW); - while (m < Mrndp) { - #else - while (m < M) { - #endif - //---------------------------------------------------- - // 'sum' accumulates the matrix A x B computation - // split across 64 lanes. - // - // YTILE represents how many column of weight matrix - // are being worked on by each wave. - //---------------------------------------------------- - float sum[N][YTILE] = {}; - scalar8 sum4[N][YTILE] = {}; - - //---------------------------------------------------- - // Fetch weight matrix B in interleaved K-split! - // - Each thread (lane) is fetching 8 elements (A_Chunk) - // - Each wave will fetch 64*8=> 512 elements (1024B) - // - YTILE represents the number of column being serviced - // by wave - // - Loop for fetching weight matrix (B) are unrolled - // - // Fetch activation matrix A from LDS - // - Loop for fetching activation matrix (A) are unrolled - // - // Finally, do the matrix multiplication in an unrolled - // fashion. This provides lot of food for compiler - // scheduling. - // - // TODO: Logic below will only work when K is multiple of 8 - //---------------------------------------------------- - for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - bigType bigA[N][UNRL] = {}; - bigType bigB[YTILE][UNRL]; - - #ifdef PCML - if ((k1 == 0) || (k1 == kBase + kFit)) { // load next chunk of A[] to LDS - if (k1 != 0) kBase += kFit; - __syncthreads(); - for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) { - uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); - if (kBase + kOff >= Kap) break; - if (kOff >= kFit) break; - for (uint32_t n = 0; n < N; n++) { - uint32_t k_in = kBase + n * Kap + kOff; - uint32_t k_ot = n * kFit + kOff; - #if defined(__gfx950__) - __builtin_amdgcn_global_load_lds((int*)(&A[k_in]), (int*)(&s[k_ot]), - 16, 0, 0); - #else - *((bigType*)(&s[k_ot])) = *((bigType*)(&A[k_in])); - #endif - } - } - __syncthreads(); - } - if (m >= M) continue; - #endif - - // Fetch the weight matrix from memory! - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; - for (int y = 0; y < YTILE; y++) - bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); - } - - // Fetch activation matrix from either just LDS or from both LDS / memory - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - for (int n = 0; n < N; n++) { - #ifdef PCML - bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); - #else - if (k_ + Kap * n < max_lds_len) - bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); - else - bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); - #endif - } - } - - // Do the matrix multiplication in interleaved manner - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - for (uint32_t n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if constexpr (!use_mfma) - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) - } - else - for (uint32_t b = 0; b < A_CHUNK / 4; b++) - sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( - bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); - } - } - } - } - - #ifdef PCML - if (m >= M) { - m += CuCount * _WvPrGrp * YTILE; - kBase = 0; - continue; - } - #endif - - //---------------------------------------------------- - // Final reduction step using shuffle - //---------------------------------------------------- - if constexpr (!use_mfma) { - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, - 1); // row_shr8 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, - 1); // row_shr4 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, - 1); // row_shr2 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, - 1); // row_shr1 - #if defined(__HIP__GFX9__) - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - #else - sum[n][y] += __shfl_xor(sum[n][y], 16); - #endif - } - } - - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if (commitColumn[y]) { - if constexpr (std::is_same_v) { - sum[n][y] += __half2float(biases[n][y]); - } else if constexpr (std::is_same_v) { - sum[n][y] += __bfloat162float(biases[n][y]); - } - C[m + y + n * M] = __float2s(sum[n][y]); - } - } - } - } - } else { - #ifdef __HIP__GFX9__ - #pragma unroll - for (int n = 0; n < N; n++) { - #pragma unroll - for (int y = 0; y < YTILE; y++) { - float accm = sum4[n][y][0]; - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, - 1); // row_shl1 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, - 1); // row_shl2 - accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, - 1); // row_shl3 - accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, - 1); // row_shl4 - accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, - 1); // row_shl8 - accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, - 1); // row_shr15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, - 1); // ROW_BCAST15 - accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, - 1); // ROW_BCAST31 - sum4[n][y][0] = accm; - } - } - if (threadIdx.x == (THRDS - 1)) { - scalar_t biases[N][YTILE] = {}; - if (BIAS) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; - } - } - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - if (commitColumn[y]) { - sum4[n][y][0] += __bfloat162float(biases[n][y]); - C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); - } - } - } - } - #endif // __HIP__GFX9__ (MFMA path) - } - - m += CuCount * _WvPrGrp * YTILE; - kBase = 0; - - // Check whether there will be fragmentation! - // This will happen only for the last wave! - if (m < M && (m + YTILE) >= M) { - uint32_t startColumn = M - YTILE; - for (uint32_t i = 0; i < (m - startColumn); i++) { - commitColumn[i] = 0; - } - m = startColumn; - } - } -} -#else -template -__global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, - const int M, const int Bx, const int By, - const scalar_t* B, - const scalar_t* __restrict__ A, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { - UNREACHABLE_CODE -} -#endif - -// Find the min val of div2 that doesn't increase N/(div1*div2) -int mindiv(int N, int div1, int div2) { - int nPrRnd = div1 * div2; - int rnds[13]; - for (int i = 0; i < 13; i++) { - rnds[i] = (N + nPrRnd - 1) / nPrRnd; - nPrRnd -= div1; - } - for (int i = 12; i >= 0; i--) - if (rnds[0] == rnds[i]) return (div2 - i); - return 0; -} - torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, const std::optional& in_bias, const int64_t CuCount) { @@ -1287,15 +36,11 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const int max_lds_len = get_lds_size() / 2; -// Like WVSPLITK_CFG but lets the caller pick A_CHUNK. The original macro -// hard-coded A_CHUNK=8; this is a compatibility-preserving expansion that -// lets a single dispatch branch select wider/narrower per-thread global -// loads. Used today by the gfx11 K=2048 N=1 fast path (Qwen3.5 decode): -// AC=16 with WvPrGrp=32 and UNRL=8 beats the (AC=8, W=16, UN=4) default -// by ~30% on M ∈ {256, 1024} and ~4% on M=248320. Reproduce the YT/UN -// half of the sweep via benchmarks/kernels/sweep_bf16_kernel.py; the -// A_CHUNK and WvPrGrp axes are not yet exposed by wvSplitK_sweep (TODO: -// extend the op to 4 axes like wvSplitK_int4g_sweep already does). + // Dispatch macros removed — dispatch logic moved to per-N shards in + // skinny_gemms/dispatch.cuh. The fused variants below still use the + // kernels directly (N=1 only, not split). + +// Kept for fused_silu_mul and fused_silu_gate_mul (N=1 only, not split): #define WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, _AC) \ { \ dim3 block(_THRDS, _WVPRGRP); \ @@ -1317,7 +62,6 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, CuCount); \ } -// Backwards-compatible wrapper: existing call sites get A_CHUNK=8. #define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \ WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, 8) @@ -1447,41 +191,26 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, : nullptr; fptype* c = reinterpret_cast(out_c.data_ptr()); - // first shoot for biggest tile-size that keeps all simd busy, - // then cut the active waves to balance their distribution... - int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); - - const bool use_wave32 = on_gfx1x(); switch (N_in) { case 1: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 1) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 1) + launch_wvsplitk_n1(grid, stream, K_in, Kap_in, Kbp_in, M_in, Bx_in, + By_in, af4, bf4, biasf4, c, CuCount, max_lds_len); break; case 2: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 2) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 2) + launch_wvsplitk_n2(grid, stream, K_in, Kap_in, Kbp_in, M_in, Bx_in, + By_in, af4, bf4, biasf4, c, CuCount, max_lds_len); break; case 3: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 3) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 3) + launch_wvsplitk_n3(grid, stream, K_in, Kap_in, Kbp_in, M_in, Bx_in, + By_in, af4, bf4, biasf4, c, CuCount, max_lds_len); break; case 4: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 4) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 4) + launch_wvsplitk_n4(grid, stream, K_in, Kap_in, Kbp_in, M_in, Bx_in, + By_in, af4, bf4, biasf4, c, CuCount, max_lds_len); break; case 5: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 5) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 5) + launch_wvsplitk_n5(grid, stream, K_in, Kap_in, Kbp_in, M_in, Bx_in, + By_in, af4, bf4, biasf4, c, CuCount, max_lds_len); break; default: throw std::runtime_error( diff --git a/csrc/rocm/skinny_gemms/dispatch.cuh b/csrc/rocm/skinny_gemms/dispatch.cuh new file mode 100644 index 000000000000..9d95ec7b7147 --- /dev/null +++ b/csrc/rocm/skinny_gemms/dispatch.cuh @@ -0,0 +1,108 @@ +#pragma once + +#include "kernel.cuh" + +// Per-N dispatch for wvSplitK (bf16/fp16 GEMM). +// Wraps WVSPLIT_TILE_CFG for a single N value. +template +inline void dispatch_wvsplitk(dim3 grid, cudaStream_t stream, int K_in, + int Kap_in, int Kbp_in, int M_in, int Bx_in, + int By_in, const scalar_t* af4, + const scalar_t* bf4, const scalar_t* biasf4, + scalar_t* c, int CuCount, int max_lds_len) { + using fptype = scalar_t; + int N_in = N_VAL; + int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); + + // clang-format off + // These macros reference local variables above and kernel templates from + // kernel.cuh. They are copies of the macros in the original skinny_gemms.cu. + +#define WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, _AC) \ + { \ + dim3 block(_THRDS, _WVPRGRP); \ + int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, _WVPRGRP); \ + if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \ + wvSplitK_hf_sml_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + else if (Kbp_in * N_in <= max_lds_len * 1.2) \ + wvSplitK_hf_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + else \ + wvSplitK_hf_big_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + } + +#define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N, 8) + +#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \ + { \ + bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ + if (is_gfx11()) { \ + if (_sYT <= 1) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if (K_in < 1024) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 4, __N) \ + else if ((K_in % 1024 == 512) && (_sYT >= 40 || K_in >= 4096)) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ + else if ((K_in == 2048) && (__N == 1)) \ + WVSPLITK_CFG_AC(32, 32, 1, 8, __N, 16) \ + else if ((K_in == 2048) && (__N == 2 || __N == 3)) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ + else if ((K_in == 4096) && (__N == 1)) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 4, __N, 32) \ + else if ((K_in == 4096) && (__N == 2) && (M_in < 4096)) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ + else if ((K_in == 4096) && (__N == 2) && (M_in >= 4096)) \ + WVSPLITK_CFG_AC(_THRDS, 32, 1, 2, __N, 32) \ + else if ((K_in == 4096) && (__N == 3)) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 2, __N, 32) \ + else if ((K_in == 8192) && (__N == 2)) \ + WVSPLITK_CFG_AC(_THRDS, 32, 1, 2, __N, 32) \ + else if ((K_in % 2048 == 0) && (__N == 2)) \ + WVSPLITK_CFG_AC(_THRDS, 32, 2, 4, __N, 16) \ + else if ((K_in % 2048 == 0) && (__N != 2)) \ + WVSPLITK_CFG_AC(_THRDS, _WVPRGRP, 1, 4, __N, 16) \ + else if (K_in <= 2048) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if (__N >= 2 && !fit_lds) { \ + if (K_in % 1024 == 0 && Kbp_in < max_lds_len / 2) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 4, __N) \ + else \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + } else if (__N == 1) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 2, __N) \ + else \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 1, __N) \ + } else { \ + if (_sYT <= 1) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \ + else if (_sYT <= 4 * 3) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \ + else if (__N == 4) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ + else \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ + } \ + } + + if (on_gfx1x()) { + WVSPLIT_TILE_CFG(32, 16, sYT, N_VAL) + } else { + WVSPLIT_TILE_CFG(64, 16, sYT, N_VAL) + } + +#undef WVSPLIT_TILE_CFG +#undef WVSPLITK_CFG_AC +#undef WVSPLITK_CFG + // clang-format on +} diff --git a/csrc/rocm/skinny_gemms/instantiate_n1.cu b/csrc/rocm/skinny_gemms/instantiate_n1.cu new file mode 100644 index 000000000000..2ecc9dfe99d5 --- /dev/null +++ b/csrc/rocm/skinny_gemms/instantiate_n1.cu @@ -0,0 +1,22 @@ +// Per-N kernel instantiation shard for wvSplitK bf16/fp16 (N=1). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_wvsplitk_n1(dim3 grid, cudaStream_t stream, int K_in, int Kap_in, + int Kbp_in, int M_in, int Bx_in, int By_in, + const scalar_t* af4, const scalar_t* bf4, + const scalar_t* biasf4, scalar_t* c, int CuCount, + int max_lds_len) { + dispatch_wvsplitk(grid, stream, K_in, Kap_in, Kbp_in, M_in, + Bx_in, By_in, af4, bf4, biasf4, c, CuCount, + max_lds_len); +} + +template void launch_wvsplitk_n1(dim3, cudaStream_t, int, int, int, int, + int, int, const half*, const half*, + const half*, half*, int, int); +template void launch_wvsplitk_n1<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, int, int, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, int, int); diff --git a/csrc/rocm/skinny_gemms/instantiate_n2.cu b/csrc/rocm/skinny_gemms/instantiate_n2.cu new file mode 100644 index 000000000000..b75df6d6c6c7 --- /dev/null +++ b/csrc/rocm/skinny_gemms/instantiate_n2.cu @@ -0,0 +1,22 @@ +// Per-N kernel instantiation shard for wvSplitK bf16/fp16 (N=2). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_wvsplitk_n2(dim3 grid, cudaStream_t stream, int K_in, int Kap_in, + int Kbp_in, int M_in, int Bx_in, int By_in, + const scalar_t* af4, const scalar_t* bf4, + const scalar_t* biasf4, scalar_t* c, int CuCount, + int max_lds_len) { + dispatch_wvsplitk(grid, stream, K_in, Kap_in, Kbp_in, M_in, + Bx_in, By_in, af4, bf4, biasf4, c, CuCount, + max_lds_len); +} + +template void launch_wvsplitk_n2(dim3, cudaStream_t, int, int, int, int, + int, int, const half*, const half*, + const half*, half*, int, int); +template void launch_wvsplitk_n2<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, int, int, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, int, int); diff --git a/csrc/rocm/skinny_gemms/instantiate_n3.cu b/csrc/rocm/skinny_gemms/instantiate_n3.cu new file mode 100644 index 000000000000..86fc47d7b664 --- /dev/null +++ b/csrc/rocm/skinny_gemms/instantiate_n3.cu @@ -0,0 +1,22 @@ +// Per-N kernel instantiation shard for wvSplitK bf16/fp16 (N=3). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_wvsplitk_n3(dim3 grid, cudaStream_t stream, int K_in, int Kap_in, + int Kbp_in, int M_in, int Bx_in, int By_in, + const scalar_t* af4, const scalar_t* bf4, + const scalar_t* biasf4, scalar_t* c, int CuCount, + int max_lds_len) { + dispatch_wvsplitk(grid, stream, K_in, Kap_in, Kbp_in, M_in, + Bx_in, By_in, af4, bf4, biasf4, c, CuCount, + max_lds_len); +} + +template void launch_wvsplitk_n3(dim3, cudaStream_t, int, int, int, int, + int, int, const half*, const half*, + const half*, half*, int, int); +template void launch_wvsplitk_n3<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, int, int, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, int, int); diff --git a/csrc/rocm/skinny_gemms/instantiate_n4.cu b/csrc/rocm/skinny_gemms/instantiate_n4.cu new file mode 100644 index 000000000000..eed52fe329cd --- /dev/null +++ b/csrc/rocm/skinny_gemms/instantiate_n4.cu @@ -0,0 +1,22 @@ +// Per-N kernel instantiation shard for wvSplitK bf16/fp16 (N=4). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_wvsplitk_n4(dim3 grid, cudaStream_t stream, int K_in, int Kap_in, + int Kbp_in, int M_in, int Bx_in, int By_in, + const scalar_t* af4, const scalar_t* bf4, + const scalar_t* biasf4, scalar_t* c, int CuCount, + int max_lds_len) { + dispatch_wvsplitk(grid, stream, K_in, Kap_in, Kbp_in, M_in, + Bx_in, By_in, af4, bf4, biasf4, c, CuCount, + max_lds_len); +} + +template void launch_wvsplitk_n4(dim3, cudaStream_t, int, int, int, int, + int, int, const half*, const half*, + const half*, half*, int, int); +template void launch_wvsplitk_n4<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, int, int, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, int, int); diff --git a/csrc/rocm/skinny_gemms/instantiate_n5.cu b/csrc/rocm/skinny_gemms/instantiate_n5.cu new file mode 100644 index 000000000000..865b9f4f5ae6 --- /dev/null +++ b/csrc/rocm/skinny_gemms/instantiate_n5.cu @@ -0,0 +1,22 @@ +// Per-N kernel instantiation shard for wvSplitK bf16/fp16 (N=5). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_wvsplitk_n5(dim3 grid, cudaStream_t stream, int K_in, int Kap_in, + int Kbp_in, int M_in, int Bx_in, int By_in, + const scalar_t* af4, const scalar_t* bf4, + const scalar_t* biasf4, scalar_t* c, int CuCount, + int max_lds_len) { + dispatch_wvsplitk(grid, stream, K_in, Kap_in, Kbp_in, M_in, + Bx_in, By_in, af4, bf4, biasf4, c, CuCount, + max_lds_len); +} + +template void launch_wvsplitk_n5(dim3, cudaStream_t, int, int, int, int, + int, int, const half*, const half*, + const half*, half*, int, int); +template void launch_wvsplitk_n5<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, int, int, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, int, int); diff --git a/csrc/rocm/skinny_gemms/kernel.cuh b/csrc/rocm/skinny_gemms/kernel.cuh new file mode 100644 index 000000000000..a294e22842c3 --- /dev/null +++ b/csrc/rocm/skinny_gemms/kernel.cuh @@ -0,0 +1,1258 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include +#include + +#include "../../cuda_compat.h" +#include "../../dispatch_utils.h" + +// TODO(rasmith): The kernels in this file are susceptible to integer overflow +// issues, do not take strides, and are unable to handle PyTorch tensors that +// return is_contiguous() as False (the tensors may actually be contiguous +// in memory). +// +// However, it may be possible to fix these kernels to handle both issues. + +#if defined(__HIPCC__) && \ + (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__GFX9__ +#endif + +// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts +#if defined(__GFX11__) || defined(__GFX12__) + #define __HIP__GFX1X__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__MI3XX__ +#endif + +#if defined(__gfx950__) + #define LDS_SIZE 160 * 1024 +#else + #define LDS_SIZE 64 * 1024 +#endif + +inline int get_lds_size() { + static const int result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx95") == std::string::npos ? 64 * 1024 + : 160 * 1024; + }(); + return result; +} + +inline bool on_gfx1x() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx11") != std::string::npos || + device_arch.find("gfx12") != std::string::npos; + }(); + return result; +} + +inline bool on_gfx12() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx12") != std::string::npos; + }(); + return result; +} + +inline bool is_gfx11() { + static const bool result = [] { + auto dprops = at::cuda::getCurrentDeviceProperties(); + std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx11") != std::string::npos; + }(); + return result; +} + +#if defined(NDEBUG) + #undef NDEBUG + #include + #define UNREACHABLE_CODE assert(false); + #define NDEBUG +#else + #define UNREACHABLE_CODE assert(false); +#endif + +template +struct scalar {}; + +template +struct scalar2 {}; + +template +__device__ __forceinline__ float2 __s22float2(T v); + +template +__device__ __forceinline__ T __float2s(float v); + +template +__device__ __forceinline__ T __float22s2_rn(float2 v); + +// Definitions and cvt functions for fp16 +template <> +struct scalar { + using type = half; +}; + +template <> +struct scalar2 { + using type = __half2; +}; + +template <> +__device__ __forceinline__ half __float2s(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ float2 __s22float2(__half2 v) { + return __half22float2(v); +} + +template <> +__device__ __forceinline__ __half2 __float22s2_rn(float2 v) { + return __float22half2_rn(v); +} + +// Definitions and cvt functions for bf16 +template <> +struct scalar { + using type = __hip_bfloat16; +}; + +template <> +struct scalar2 { + using type = __hip_bfloat162; +}; + +template <> +__device__ __forceinline__ __hip_bfloat16 __float2s(float v) { + return __float2bfloat16(v); +} + +template <> +__device__ __forceinline__ float2 __s22float2(__hip_bfloat162 v) { + return __bfloat1622float2(v); +} + +template <> +__device__ __forceinline__ __hip_bfloat162 __float22s2_rn(float2 v) { + return __float22bfloat162_rn(v); +} + +template +__device__ __forceinline__ T loadnt(T* addr) { + return __builtin_nontemporal_load(addr); +} + +__device__ __forceinline__ float4 load_ntmprl(const float4* addr) { + auto addr_alias = reinterpret_cast(addr); + auto dat0 = loadnt(addr_alias); + auto dat1 = loadnt(addr_alias + 1); + auto dat2 = loadnt(addr_alias + 2); + auto dat3 = loadnt(addr_alias + 3); + return make_float4(dat0, dat1, dat2, dat3); +} + +// LLGemm1 kernel + LLMM1 entry point — only compiled in the main TU +// (skinny_gemms.cu), not in per-N instantiation shards. +#ifdef SKINNY_GEMMS_MAIN_TU +template +__global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b, + scalar_t* out_c, const int K) { + using scalar2_t = typename scalar2::type; + auto af4 = reinterpret_cast(in_a); + auto bf4 = reinterpret_cast(in_b); + auto c = reinterpret_cast(out_c); + __shared__ float red_smem[NUM_A_ROWS_PER_BLOCK][WARP_SIZE]; + const int row_addr = blockIdx.x * NUM_A_ROWS_PER_BLOCK * K / 8; + const int threadid = threadIdx.x; + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int num_warps = blockDim.x / WARP_SIZE; + const int qwarpid = threadid / 16; + const int qthreadid = threadid % 16; + float4 rowA_elem4[NUM_A_ROWS_PER_BLOCK]; + scalar2_t colB_elem4x, colB_elem4y, colB_elem4z, colB_elem4w; + float acc[NUM_A_ROWS_PER_BLOCK]; + scalar2_t acch2; + scalar2_t oval; + + // Each thread processes 8 elements per iteration. With NUM_THREADS threads, + // each iteration covers NUM_THREADS * 8 elements of K. Loop over chunks + // to handle K values larger than NUM_THREADS * 8. + const int elems_per_iter = blockDim.x; // threads, each handling 8 elements + const int K_div8 = K / 8; + + #pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + acc[i] = 0.f; + } + + for (int base = 0; base < K_div8; base += elems_per_iter) { + int idx = base + threadid; + if (idx < K_div8) { + #pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + rowA_elem4[i] = load_ntmprl(&af4[row_addr + idx + K_div8 * i]); + } + colB_elem4x = bf4[idx * 4 + 0]; + colB_elem4y = bf4[idx * 4 + 1]; + colB_elem4z = bf4[idx * 4 + 2]; + colB_elem4w = bf4[idx * 4 + 3]; + + scalar2_t Af2; + auto Ah2ptr = reinterpret_cast(&rowA_elem4); + scalar2_t* ah2lptr; + + #pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + ah2lptr = Ah2ptr + i * 4; + Af2 = *(ah2lptr); + acch2 = __hmul2(Af2, colB_elem4x); + Af2 = *(ah2lptr + 1); + acch2 = __hfma2(Af2, colB_elem4y, acch2); + Af2 = *(ah2lptr + 2); + acch2 = __hfma2(Af2, colB_elem4z, acch2); + Af2 = *(ah2lptr + 3); + acch2 = __hfma2(Af2, colB_elem4w, acch2); + float2 S = __s22float2(acch2); + acc[i] += S.x + S.y; + } + } + } + + // all reduce across warp. + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + #pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + acc[i] += __shfl_xor(acc[i], mask); + } + } + + // Warp leaders store the data to shared memory. + if (lane < NUM_A_ROWS_PER_BLOCK) { + red_smem[lane][warp] = acc[lane]; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + if (qwarpid < NUM_A_ROWS_PER_BLOCK) { + acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f; + #pragma unroll + for (int mask = 16 / 2; mask >= 1; mask /= 2) { + acc[qwarpid] += __shfl_xor(acc[qwarpid], mask); + } + float oval2 = __shfl_xor(acc[qwarpid], 16); + + if (lane % 32 == 0) { + oval = __float22s2_rn(make_float2(acc[qwarpid], oval2)); + c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval; + } + } +} + +torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, + const int64_t rows_per_block) { + auto M = in_a.size(0); + auto K = in_a.size(1); + auto N = in_b.size(0); + + TORCH_CHECK(N == 1, "Row number of activation tensor must be 1."); + TORCH_CHECK(in_a.dtype() == in_b.dtype()); + TORCH_CHECK(in_b.dtype() == torch::kFloat16 || + in_b.dtype() == torch::kBFloat16); + + auto out_c = torch::empty( + {N, M}, torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device())); + + // NUM_THREADS must be a multiple of WARP_SIZE (warp shuffle operations). + // Cap at 512 (16 warps) because the cross-warp reduction uses 16 threads + // per row. The kernel loops over K in chunks when K/8 > NUM_THREADS. + int NUM_THREADS = + max(rows_per_block * 16, + K * 2 / 16 % WARP_SIZE == 0 + ? K * 2 / 16 + : K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE)); + NUM_THREADS = min(NUM_THREADS, 512); + + int NUM_BLOCKS = M / rows_per_block; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_b)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // call the kernel function... + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "LLGemm1", [&] { + auto a_ptr = in_a.data_ptr(); + auto b_ptr = in_b.data_ptr(); + auto c_ptr = out_c.data_ptr(); + if (rows_per_block == 2) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 4) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 8) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 16) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else { + NUM_BLOCKS = M / 4; + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } + }); + + return out_c; +} +#endif // SKINNY_GEMMS_MAIN_TU + +#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__) + #define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + asm("v_dot2c_f32_f16 %0, %2, %3" \ + : "=v"(V0) \ + : "0"(V0), "v"(V2), "v"(V3)); \ + } else if constexpr (std::is_same_v) { \ + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ + V0 += (s.x + s.y); \ + } +#elif defined(__HIP__GFX1X__) + // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12) + #define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ + } else if constexpr (std::is_same_v) { \ + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ + V0 += (s.x + s.y); \ + } +#endif + +// To avoid LLVM silently upcasting to double +__device__ inline unsigned int min__(uint32_t a, uint32_t b) { + return min(a, b); +} + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// META3-2: Variant of the wvSplitK_hf_sml LDS-load loop that fuses the +// silu_and_mul preamble. The source activation tensor A has 2*K columns +// per row, packed as [gate(K) | up(K)]; the LDS staging buffer ends up +// with silu(gate) * up. N=1 only (decode batch=1 path). Mirrors the +// int4 EXPERIMENT-g helper `load_act_into_lds_silu_mul` for the bf16 / +// fp16 wvSplitK template. +// +// Match the unfused silu_and_mul semantics exactly: silu in fp32, cast +// back to scalar_t, then the scalar_t multiply by `up`. Doing the +// multiply in fp32 changed downstream GEMM rounding subtly enough to +// regress generated text on the int4 path; same caution applies here. +template +__device__ __forceinline__ void load_act_into_lds_silu_mul_bf16( + scalar_t* s, const scalar_t* __restrict__ A, const int K, + const int max_lds_len) { + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + scalar8 h8; + }; + const int limit = min__(K, max_lds_len); + for (uint32_t k = 0; k < (uint32_t)limit; k += THRDS * WvPrGrp * A_CHUNK) { + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); + if (k_in >= (uint32_t)limit) break; + bigType gate = *((const bigType*)(&A[k_in])); + bigType up = *((const bigType*)(&A[k_in + K])); + bigType out; + #pragma unroll + for (int i = 0; i < A_CHUNK; ++i) { + float g; + if constexpr (std::is_same_v) { + g = __bfloat162float(gate.h[i]); + } else { + g = __half2float(gate.h[i]); + } + scalar_t silu_g = __float2s(g / (1.0f + expf(-g))); + out.h[i] = silu_g * up.h[i]; + } + *((bigType*)(&s[k_in])) = out; + } + __syncthreads(); +} + +// This version targets cases where A[] fits LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount, + const scalar_t* __restrict__ GATE = nullptr) { + static_assert(!FUSED_SILU_MUL || N == 1, + "FUSED_SILU_MUL is only supported with N=1"); + static_assert(!FUSED_GATE_MUL || N == 1, + "FUSED_GATE_MUL is only supported with N=1"); + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + //---------------------------------------------------- + // Reserving 64/160 KB of LDS to have 1 WG / CU + // Goal is to bring the activation matrix A to the LDS + // and use it across the lifetime of the work group + // TODO: When activation matrix is larger than 64 KB + // then this is not going to work! + //---------------------------------------------------- + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Fetch the activation matrix to LDS + // Loop iteration: + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements + // - Each WG will fetch 512 * 16 => 8K elements + // - Then the WG will move to another 8 K elements + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + if constexpr (FUSED_SILU_MUL) { + // META3-2: A is laid out [gate(K) | up(K)] per row; this writes + // silu(gate)*up into LDS so the kernel sees the post-activation K + // elements directly. N=1 (asserted above). + load_act_into_lds_silu_mul_bf16( + s, A, K, max_lds_len); + } else { + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + __syncthreads(); + } + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; + + //---------------------------------------------------- + // Each wave works on a single column of weight matrix. + // There are 16 waves per WG, and hence, each WG is + // working on 16 columns of weight matrix. Moreover, + // we tile in column direction by YTILE, so when YTILE=1 + // the above math is right, however, when YTILE=2 then + // each wave will be working on 2 columns and WG will + // be working on 32 columns. + // + // Top level loop that makes WGs persistent! + // - WGs iterates across columns of weight matrix + // - Each wave within WG works on a given column(s) + // - After completing first set of columns, WGs start + // working on the next set of available columns + //---------------------------------------------------- + while (m < M) { + //---------------------------------------------------- + // 'sum' accumulates the matrix A x B computation + // split across 64 lanes. + // + // YTILE represents how many column of weight matrix + // are being worked on by each wave. + //---------------------------------------------------- + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + __builtin_amdgcn_sched_barrier(0); + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + scalar_t out_val = __float2s(sum[n][y]); + if constexpr (FUSED_GATE_MUL) { + // META3-2 Phase 2: per-token scalar mul (mirrors the + // unfused `F.sigmoid(expert_gate(x)) * out` in scalar_t). + out_val = out_val * GATE[n]; + } + C[m + y + n * M] = out_val; + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + /*float accm1 = 0; + for (int i=0; i<64; i++) + accm1 += __shfl(sum4[n][y][i%4], i); + sum4[n][y][0] = accm1;*/ + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + scalar_t out_val = __float2bfloat16(sum4[n][y][0]); + if constexpr (FUSED_GATE_MUL) { + // META3-2 Phase 2: per-token scalar mul. + out_val = out_val * GATE[n]; + } + C[m + y + n * M] = out_val; + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + m += CuCount * _WvPrGrp * YTILE; + } +} +#else +template +__global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount, + const scalar_t* __restrict__ GATE = nullptr) { + UNREACHABLE_CODE +} +#endif + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// This version targets cases where A[] marginally exceeds LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Computation of columns that need to be committed to memory! + //---------------------------------------------------- + uint32_t commitColumn[YTILE]; + for (uint32_t i = 0; i < YTILE; i++) { + commitColumn[i] = 1; + } + + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + while (m < M) { + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + for (uint32_t n = 0; n < N; n++) { + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(sum[n][y]); + } + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + // float accm1 = 0; + // for (int i=0; i<64; i++) + // accm1 += __shfl(sum4[n][y][i%4], i); + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); + } + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + + m += CuCount * _WvPrGrp * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + } +} + +#else +template +__global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// This version targets big A[] cases, where it is much larger than LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = __attribute__((__vector_size__(4 * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + //---------------------------------------------------- + // Reserving 64/160 KB of LDS to have 1 WG / CU + // Goal is to bring the activation matrix A to the LDS + // and use it across the lifetime of the work group + // TODO: When activation matrix is larger than 64 KB + // then this is not going to work! + //---------------------------------------------------- + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Computation of columns that need to be committed to memory! + //---------------------------------------------------- + uint32_t commitColumn[YTILE]; + for (uint32_t i = 0; i < YTILE; i++) { + commitColumn[i] = 1; + } + + // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp); + if (threadIdx.y >= _WvPrGrp) return; + + //---------------------------------------------------- + // Indexing function into the column of weight matrix B + // Algorithm does 64 lane k-splitting / wave and uses + // WG ID and Thread ID to find the index. + //---------------------------------------------------- + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + + //---------------------------------------------------- + // Fetch the activation matrix to LDS + // Loop iteration: + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements + // - Each WG will fetch 512 * 16 => 8K elements + // - Then the WG will move to another 8 K elements + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + #define PCML + #ifndef PCML + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + __syncthreads(); + #endif + + #define TUC (THRDS * UNRL * A_CHUNK) + uint32_t kBase = 0; + // find biggest k size that fits in LDS + uint32_t kFit = (max_lds_len) / N; + // kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple + // of TUC + kFit = (kFit % TUC == 0) + ? kFit + : (kFit - kFit % TUC); // round up to multiple of TUC + // if (kFit == 0) kFit = TUC; + kFit = min__(kFit, Kap); + + //---------------------------------------------------- + // Each wave works on a single column of weight matrix. + // There are 16 waves per WG, and hence, each WG is + // working on 16 columns of weight matrix. Moreover, + // we tile in column direction by YTILE, so when YTILE=1 + // the above math is right, however, when YTILE=2 then + // each wave will be working on 2 columns and WG will + // be working on 32 columns. + // + // Top level loop that makes WGs persistent! + // - WGs iterates across columns of weight matrix + // - Each wave within WG works on a given column(s) + // - After completing first set of columns, WGs start + // working on the next set of available columns + //---------------------------------------------------- + #ifdef PCML + int YW = (YTILE * _WvPrGrp); + uint32_t Mrndp = (M % YW == 0) ? M : (M - M % YW + YW); + while (m < Mrndp) { + #else + while (m < M) { + #endif + //---------------------------------------------------- + // 'sum' accumulates the matrix A x B computation + // split across 64 lanes. + // + // YTILE represents how many column of weight matrix + // are being worked on by each wave. + //---------------------------------------------------- + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + //---------------------------------------------------- + // Fetch weight matrix B in interleaved K-split! + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements (1024B) + // - YTILE represents the number of column being serviced + // by wave + // - Loop for fetching weight matrix (B) are unrolled + // + // Fetch activation matrix A from LDS + // - Loop for fetching activation matrix (A) are unrolled + // + // Finally, do the matrix multiplication in an unrolled + // fashion. This provides lot of food for compiler + // scheduling. + // + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + + #ifdef PCML + if ((k1 == 0) || (k1 == kBase + kFit)) { // load next chunk of A[] to LDS + if (k1 != 0) kBase += kFit; + __syncthreads(); + for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) { + uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); + if (kBase + kOff >= Kap) break; + if (kOff >= kFit) break; + for (uint32_t n = 0; n < N; n++) { + uint32_t k_in = kBase + n * Kap + kOff; + uint32_t k_ot = n * kFit + kOff; + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k_in]), (int*)(&s[k_ot]), + 16, 0, 0); + #else + *((bigType*)(&s[k_ot])) = *((bigType*)(&A[k_in])); + #endif + } + } + __syncthreads(); + } + if (m >= M) continue; + #endif + + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + #ifdef PCML + bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); + #else + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + #endif + } + } + + // Do the matrix multiplication in interleaved manner + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + + #ifdef PCML + if (m >= M) { + m += CuCount * _WvPrGrp * YTILE; + kBase = 0; + continue; + } + #endif + + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(sum[n][y]); + } + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); + } + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + + m += CuCount * _WvPrGrp * YTILE; + kBase = 0; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + } +} +#else +template +__global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif + +// Find the min val of div2 that doesn't increase N/(div1*div2) +inline int mindiv(int N, int div1, int div2) { + int nPrRnd = div1 * div2; + int rnds[13]; + for (int i = 0; i < 13; i++) { + rnds[i] = (N + nPrRnd - 1) / nPrRnd; + nPrRnd -= div1; + } + for (int i = 12; i >= 0; i--) + if (rnds[0] == rnds[i]) return (div2 - i); + return 0; +} diff --git a/csrc/rocm/skinny_gemms/launch.h b/csrc/rocm/skinny_gemms/launch.h new file mode 100644 index 000000000000..1bccf558bd22 --- /dev/null +++ b/csrc/rocm/skinny_gemms/launch.h @@ -0,0 +1,21 @@ +#pragma once + +#include +#include +#include + +#define DECLARE_LAUNCH_BF16(N) \ + template \ + void launch_wvsplitk_n##N(dim3 grid, cudaStream_t stream, int K_in, \ + int Kap_in, int Kbp_in, int M_in, int Bx_in, \ + int By_in, const scalar_t* af4, \ + const scalar_t* bf4, const scalar_t* biasf4, \ + scalar_t* c, int CuCount, int max_lds_len); + +DECLARE_LAUNCH_BF16(1) +DECLARE_LAUNCH_BF16(2) +DECLARE_LAUNCH_BF16(3) +DECLARE_LAUNCH_BF16(4) +DECLARE_LAUNCH_BF16(5) + +#undef DECLARE_LAUNCH_BF16 diff --git a/csrc/rocm/skinny_gemms_int4.cu b/csrc/rocm/skinny_gemms_int4.cu index 9ba9c5aac751..6dfaefc7c44d 100644 --- a/csrc/rocm/skinny_gemms_int4.cu +++ b/csrc/rocm/skinny_gemms_int4.cu @@ -1,8 +1,14 @@ -// Production wrappers for int4 wvSplitK GEMMs. Templates and macros live in -// skinny_gemms_int4_kernels.cuh; the sweep variants live in -// skinny_gemms_int4_sweep.cu. Splitting kept the file small so production + -// sweep TUs compile in parallel. +// Production wrappers for int4 wvSplitK GEMMs. Kernel templates live in +// skinny_gemms_int4_kernels.cuh; per-N instantiation shards live in +// skinny_gemms_int4/instantiate_n{1..5}.cu. +#include +#include +#include + +#include "../cuda_compat.h" +#include "dispatch_utils.h" #include "skinny_gemms_int4_kernels.cuh" +#include "skinny_gemms_int4/launch.h" torch::Tensor wvSplitK_int4_g(const at::Tensor& in_w, const at::Tensor& in_x, const at::Tensor& in_scale, @@ -71,6 +77,8 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_w, const at::Tensor& in_x, const at::cuda::OptionalCUDAGuard device_guard(device_of(in_w)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + bool has_zp = in_zero_points.has_value(); + AT_DISPATCH_REDUCED_FLOATING_TYPES( in_x.scalar_type(), "wvSplitK_int4_g", [&] { using fptype = typename scalar::type; @@ -88,30 +96,46 @@ torch::Tensor wvSplitK_int4_g(const at::Tensor& in_w, const at::Tensor& in_x, : nullptr; fptype* cptr = reinterpret_cast(out_c.data_ptr()); - if (in_zero_points.has_value()) - WVSPLIT_INT4G_DISPATCH(true) - else - WVSPLIT_INT4G_DISPATCH(false) + switch (N_in) { + case 1: + launch_int4_n1(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, + sptr, zpptr, biasptr, cptr, CuCount, + b_row_stride_bytes_i32, group_size, max_lds_len, + has_zp); + break; + case 2: + launch_int4_n2(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, + sptr, zpptr, biasptr, cptr, CuCount, + b_row_stride_bytes_i32, group_size, max_lds_len, + has_zp); + break; + case 3: + launch_int4_n3(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, + sptr, zpptr, biasptr, cptr, CuCount, + b_row_stride_bytes_i32, group_size, max_lds_len, + has_zp); + break; + case 4: + launch_int4_n4(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, + sptr, zpptr, biasptr, cptr, CuCount, + b_row_stride_bytes_i32, group_size, max_lds_len, + has_zp); + break; + case 5: + launch_int4_n5(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, + sptr, zpptr, biasptr, cptr, CuCount, + b_row_stride_bytes_i32, group_size, max_lds_len, + has_zp); + break; + default: + throw std::runtime_error("Unsupported N value: " + + std::to_string(N_in)); + } }); return out_c; } -// Fused MoE wrapper around wvSplitK_int4_g. -// -// Single GPU kernel launch — expert routing happens on-device via blockIdx.y. -// No host-side loop, no GPU→CPU memcpy of expert_ids. -// Activations must be pre-permuted into contiguous expert blocks. -// -// a: [num_slots, K] pre-permuted activations (fp16/bf16) -// w: [E, N_weight, K//8] int32 packed weights (skinny layout) -// scales: [E, N_weight, K//group_size] fp16/bf16 -// c: [num_slots, N_weight] output (pre-allocated) -// expert_ids: [num_expert_blocks] int32 — expert id per block -// block_size_m: 1, 2, or 4 — rows per expert block -// CuCount: number of compute units -// group_size: 32 or 128 -// zero_points: [E, N_weight, K//group_size] or empty tensor void fused_moe_wvSplitK_int4_gemm(torch::Tensor a, torch::Tensor w, torch::Tensor scales, torch::Tensor c, torch::Tensor expert_ids, @@ -122,36 +146,25 @@ void fused_moe_wvSplitK_int4_gemm(torch::Tensor a, torch::Tensor w, const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - // Weight layout: [E, N_weight, K//8] - int M_in = static_cast(w.size(1)); // N_weight (wvSplitK M dim) - int K_in = static_cast(w.size(2)) * 8; // unpacked K - int N_in = static_cast(block_size_m); // batch rows per expert block + int M_in = static_cast(w.size(1)); + int K_in = static_cast(w.size(2)) * 8; + int N_in = static_cast(block_size_m); int num_expert_blocks = static_cast(expert_ids.size(0)); bool has_zp = zero_points.numel() > 0; - // Expert strides: w stride is in int32 elements, convert to bytes for uint8* long expert_stride_w = w.stride(0) * static_cast(sizeof(int32_t)); long expert_stride_s = scales.stride(0); long expert_stride_zp = has_zp ? zero_points.stride(0) : 0; const int max_lds_len = get_lds_size_int4() / 2; - // Scattered mode: sorted_token_ids is non-empty, kernel indexes into - // unpermuted activations via sorted_token_ids[block] / top_k. bool scattered = sorted_token_ids.numel() > 0; int top_k_in = scattered ? static_cast(top_k) : 1; - // The MOE_WVSPLIT_INT4G_GS_W_AC dispatch macro (in _kernels.cuh) takes a - // runtime fuse_silu_mul branch reachable from the sweep wrapper. The - // production op never requests fusion (its public signature has no such - // arg); declare a const-false here so the dispatch falls through to the - // unfused codepath and the optimiser eliminates the fused branch. const bool fuse_silu_mul = false; - // No c.zero_() needed: the wvSplitK kernel writes all M output rows directly - // (no atomicAdd), and padding blocks with expert_id==-1 are never read by - // the caller (moe_unpermute only accesses valid token slots). + dim3 grid(CuCount); AT_DISPATCH_REDUCED_FLOATING_TYPES( a.scalar_type(), "fused_moe_wvSplitK_int4_gemm", [&] { @@ -168,12 +181,45 @@ void fused_moe_wvSplitK_int4_gemm(torch::Tensor a, torch::Tensor w, const int* stidptr = scattered ? sorted_token_ids.data_ptr() : nullptr; - // Single kernel launch: grid = dim3(CuCount); the expert-block - // dimension is walked by an in-kernel for-loop inside the MoE - // kernel so the "workgroups == CuCount" M-split invariant holds. - if (has_zp) - MOE_WVSPLIT_INT4G_DISPATCH(true) - else - MOE_WVSPLIT_INT4G_DISPATCH(false) + switch (N_in) { + case 1: + launch_moe_int4_n1(grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, + zpptr, cptr, eidptr, stidptr, top_k_in, + expert_stride_w, expert_stride_s, + expert_stride_zp, CuCount, num_expert_blocks, + group_size, max_lds_len, has_zp, fuse_silu_mul); + break; + case 2: + launch_moe_int4_n2(grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, + zpptr, cptr, eidptr, stidptr, top_k_in, + expert_stride_w, expert_stride_s, + expert_stride_zp, CuCount, num_expert_blocks, + group_size, max_lds_len, has_zp, fuse_silu_mul); + break; + case 3: + launch_moe_int4_n3(grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, + zpptr, cptr, eidptr, stidptr, top_k_in, + expert_stride_w, expert_stride_s, + expert_stride_zp, CuCount, num_expert_blocks, + group_size, max_lds_len, has_zp, fuse_silu_mul); + break; + case 4: + launch_moe_int4_n4(grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, + zpptr, cptr, eidptr, stidptr, top_k_in, + expert_stride_w, expert_stride_s, + expert_stride_zp, CuCount, num_expert_blocks, + group_size, max_lds_len, has_zp, fuse_silu_mul); + break; + case 5: + launch_moe_int4_n5(grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, + zpptr, cptr, eidptr, stidptr, top_k_in, + expert_stride_w, expert_stride_s, + expert_stride_zp, CuCount, num_expert_blocks, + group_size, max_lds_len, has_zp, fuse_silu_mul); + break; + default: + throw std::runtime_error("Unsupported N value: " + + std::to_string(N_in)); + } }); } diff --git a/csrc/rocm/skinny_gemms_int4/dispatch.cuh b/csrc/rocm/skinny_gemms_int4/dispatch.cuh new file mode 100644 index 000000000000..d6a7b93aa355 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/dispatch.cuh @@ -0,0 +1,45 @@ +#pragma once + +#include "../skinny_gemms_int4_kernels.cuh" + +// Per-N dispatch for wvSplitK_int4_g (standard GEMM). +// Wraps WVSPLIT_INT4G_TILE for a single N value. +template +inline void dispatch_int4_g(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, + int b_row_stride_bytes_i32, int64_t group_size, + int max_lds_len, bool has_zp) { + using fptype = scalar_t; + int N_in = N_VAL; + int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); + if (has_zp) { + WVSPLIT_INT4G_TILE(sYT, N_VAL, true) + } else { + WVSPLIT_INT4G_TILE(sYT, N_VAL, false) + } +} + +// Per-N dispatch for fused_moe_wvSplitK_int4_gemm (MoE GEMM). +// Wraps MOE_WVSPLIT_INT4G_TILE for a single N value. +template +inline void dispatch_moe_int4_g(dim3 grid, cudaStream_t stream, int K_in, + int M_in, int N_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, scalar_t* cptr, + const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, + int64_t group_size, int max_lds_len, + bool has_zp, bool fuse_silu_mul) { + using fptype = scalar_t; + int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); + if (has_zp) { + MOE_WVSPLIT_INT4G_TILE(sYT, N_VAL, true) + } else { + MOE_WVSPLIT_INT4G_TILE(sYT, N_VAL, false) + } +} diff --git a/csrc/rocm/skinny_gemms_int4/instantiate_n1.cu b/csrc/rocm/skinny_gemms_int4/instantiate_n1.cu new file mode 100644 index 000000000000..6a17f0ade080 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/instantiate_n1.cu @@ -0,0 +1,51 @@ +// Per-N kernel instantiation shard for wvSplitK_int4 (N=1). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int4_n1(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, + int64_t group_size, int max_lds_len, bool has_zp) { + dispatch_int4_g( + grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, + cptr, CuCount, b_row_stride_bytes_i32, group_size, max_lds_len, has_zp); +} + +template void launch_int4_n1(dim3, cudaStream_t, int, int, int, int, + const uint8_t*, const half*, const half*, + const half*, const half*, half*, int, int, + int64_t, int, bool); +template void launch_int4_n1<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const uint8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + const __hip_bfloat16*, __hip_bfloat16*, int, int, int64_t, int, bool); + +template +void launch_moe_int4_n1(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int N_in, const uint8_t* wptr, const scalar_t* aptr, + const scalar_t* sptr, const scalar_t* zpptr, + scalar_t* cptr, const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, int64_t group_size, + int max_lds_len, bool has_zp, bool fuse_silu_mul) { + dispatch_moe_int4_g( + grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, zpptr, cptr, eidptr, + stidptr, top_k_in, expert_stride_w, expert_stride_s, expert_stride_zp, + CuCount, num_expert_blocks, group_size, max_lds_len, has_zp, + fuse_silu_mul); +} + +template void launch_moe_int4_n1(dim3, cudaStream_t, int, int, int, + const uint8_t*, const half*, const half*, + const half*, half*, const int*, + const int*, int, long, long, long, int, + int, int64_t, int, bool, bool); +template void launch_moe_int4_n1<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, const uint8_t*, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, const int*, + const int*, int, long, long, long, int, int, int64_t, int, bool, bool); diff --git a/csrc/rocm/skinny_gemms_int4/instantiate_n2.cu b/csrc/rocm/skinny_gemms_int4/instantiate_n2.cu new file mode 100644 index 000000000000..c68442ad8c9d --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/instantiate_n2.cu @@ -0,0 +1,51 @@ +// Per-N kernel instantiation shard for wvSplitK_int4 (N=2). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int4_n2(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, + int64_t group_size, int max_lds_len, bool has_zp) { + dispatch_int4_g( + grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, + cptr, CuCount, b_row_stride_bytes_i32, group_size, max_lds_len, has_zp); +} + +template void launch_int4_n2(dim3, cudaStream_t, int, int, int, int, + const uint8_t*, const half*, const half*, + const half*, const half*, half*, int, int, + int64_t, int, bool); +template void launch_int4_n2<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const uint8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + const __hip_bfloat16*, __hip_bfloat16*, int, int, int64_t, int, bool); + +template +void launch_moe_int4_n2(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int N_in, const uint8_t* wptr, const scalar_t* aptr, + const scalar_t* sptr, const scalar_t* zpptr, + scalar_t* cptr, const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, int64_t group_size, + int max_lds_len, bool has_zp, bool fuse_silu_mul) { + dispatch_moe_int4_g( + grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, zpptr, cptr, eidptr, + stidptr, top_k_in, expert_stride_w, expert_stride_s, expert_stride_zp, + CuCount, num_expert_blocks, group_size, max_lds_len, has_zp, + fuse_silu_mul); +} + +template void launch_moe_int4_n2(dim3, cudaStream_t, int, int, int, + const uint8_t*, const half*, const half*, + const half*, half*, const int*, + const int*, int, long, long, long, int, + int, int64_t, int, bool, bool); +template void launch_moe_int4_n2<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, const uint8_t*, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, const int*, + const int*, int, long, long, long, int, int, int64_t, int, bool, bool); diff --git a/csrc/rocm/skinny_gemms_int4/instantiate_n3.cu b/csrc/rocm/skinny_gemms_int4/instantiate_n3.cu new file mode 100644 index 000000000000..22561841af23 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/instantiate_n3.cu @@ -0,0 +1,51 @@ +// Per-N kernel instantiation shard for wvSplitK_int4 (N=3). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int4_n3(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, + int64_t group_size, int max_lds_len, bool has_zp) { + dispatch_int4_g( + grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, + cptr, CuCount, b_row_stride_bytes_i32, group_size, max_lds_len, has_zp); +} + +template void launch_int4_n3(dim3, cudaStream_t, int, int, int, int, + const uint8_t*, const half*, const half*, + const half*, const half*, half*, int, int, + int64_t, int, bool); +template void launch_int4_n3<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const uint8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + const __hip_bfloat16*, __hip_bfloat16*, int, int, int64_t, int, bool); + +template +void launch_moe_int4_n3(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int N_in, const uint8_t* wptr, const scalar_t* aptr, + const scalar_t* sptr, const scalar_t* zpptr, + scalar_t* cptr, const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, int64_t group_size, + int max_lds_len, bool has_zp, bool fuse_silu_mul) { + dispatch_moe_int4_g( + grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, zpptr, cptr, eidptr, + stidptr, top_k_in, expert_stride_w, expert_stride_s, expert_stride_zp, + CuCount, num_expert_blocks, group_size, max_lds_len, has_zp, + fuse_silu_mul); +} + +template void launch_moe_int4_n3(dim3, cudaStream_t, int, int, int, + const uint8_t*, const half*, const half*, + const half*, half*, const int*, + const int*, int, long, long, long, int, + int, int64_t, int, bool, bool); +template void launch_moe_int4_n3<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, const uint8_t*, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, const int*, + const int*, int, long, long, long, int, int, int64_t, int, bool, bool); diff --git a/csrc/rocm/skinny_gemms_int4/instantiate_n4.cu b/csrc/rocm/skinny_gemms_int4/instantiate_n4.cu new file mode 100644 index 000000000000..2a9fa76edc80 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/instantiate_n4.cu @@ -0,0 +1,51 @@ +// Per-N kernel instantiation shard for wvSplitK_int4 (N=4). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int4_n4(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, + int64_t group_size, int max_lds_len, bool has_zp) { + dispatch_int4_g( + grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, + cptr, CuCount, b_row_stride_bytes_i32, group_size, max_lds_len, has_zp); +} + +template void launch_int4_n4(dim3, cudaStream_t, int, int, int, int, + const uint8_t*, const half*, const half*, + const half*, const half*, half*, int, int, + int64_t, int, bool); +template void launch_int4_n4<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const uint8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + const __hip_bfloat16*, __hip_bfloat16*, int, int, int64_t, int, bool); + +template +void launch_moe_int4_n4(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int N_in, const uint8_t* wptr, const scalar_t* aptr, + const scalar_t* sptr, const scalar_t* zpptr, + scalar_t* cptr, const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, int64_t group_size, + int max_lds_len, bool has_zp, bool fuse_silu_mul) { + dispatch_moe_int4_g( + grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, zpptr, cptr, eidptr, + stidptr, top_k_in, expert_stride_w, expert_stride_s, expert_stride_zp, + CuCount, num_expert_blocks, group_size, max_lds_len, has_zp, + fuse_silu_mul); +} + +template void launch_moe_int4_n4(dim3, cudaStream_t, int, int, int, + const uint8_t*, const half*, const half*, + const half*, half*, const int*, + const int*, int, long, long, long, int, + int, int64_t, int, bool, bool); +template void launch_moe_int4_n4<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, const uint8_t*, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, const int*, + const int*, int, long, long, long, int, int, int64_t, int, bool, bool); diff --git a/csrc/rocm/skinny_gemms_int4/instantiate_n5.cu b/csrc/rocm/skinny_gemms_int4/instantiate_n5.cu new file mode 100644 index 000000000000..d6d48232819f --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/instantiate_n5.cu @@ -0,0 +1,51 @@ +// Per-N kernel instantiation shard for wvSplitK_int4 (N=5). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int4_n5(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int Bx_in, int By_in, const uint8_t* wptr, + const scalar_t* aptr, const scalar_t* sptr, + const scalar_t* zpptr, const scalar_t* biasptr, + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, + int64_t group_size, int max_lds_len, bool has_zp) { + dispatch_int4_g( + grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, zpptr, biasptr, + cptr, CuCount, b_row_stride_bytes_i32, group_size, max_lds_len, has_zp); +} + +template void launch_int4_n5(dim3, cudaStream_t, int, int, int, int, + const uint8_t*, const half*, const half*, + const half*, const half*, half*, int, int, + int64_t, int, bool); +template void launch_int4_n5<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const uint8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + const __hip_bfloat16*, __hip_bfloat16*, int, int, int64_t, int, bool); + +template +void launch_moe_int4_n5(dim3 grid, cudaStream_t stream, int K_in, int M_in, + int N_in, const uint8_t* wptr, const scalar_t* aptr, + const scalar_t* sptr, const scalar_t* zpptr, + scalar_t* cptr, const int* eidptr, const int* stidptr, + int top_k_in, long expert_stride_w, + long expert_stride_s, long expert_stride_zp, + int CuCount, int num_expert_blocks, int64_t group_size, + int max_lds_len, bool has_zp, bool fuse_silu_mul) { + dispatch_moe_int4_g( + grid, stream, K_in, M_in, N_in, wptr, aptr, sptr, zpptr, cptr, eidptr, + stidptr, top_k_in, expert_stride_w, expert_stride_s, expert_stride_zp, + CuCount, num_expert_blocks, group_size, max_lds_len, has_zp, + fuse_silu_mul); +} + +template void launch_moe_int4_n5(dim3, cudaStream_t, int, int, int, + const uint8_t*, const half*, const half*, + const half*, half*, const int*, + const int*, int, long, long, long, int, + int, int64_t, int, bool, bool); +template void launch_moe_int4_n5<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, const uint8_t*, const __hip_bfloat16*, + const __hip_bfloat16*, const __hip_bfloat16*, __hip_bfloat16*, const int*, + const int*, int, long, long, long, int, int, int64_t, int, bool, bool); diff --git a/csrc/rocm/skinny_gemms_int4/launch.h b/csrc/rocm/skinny_gemms_int4/launch.h new file mode 100644 index 000000000000..1758823d9469 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int4/launch.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include + +#define DECLARE_LAUNCH_INT4(N) \ + template \ + void launch_int4_n##N( \ + dim3 grid, cudaStream_t stream, int K_in, int M_in, int Bx_in, \ + int By_in, const uint8_t* wptr, const scalar_t* aptr, \ + const scalar_t* sptr, const scalar_t* zpptr, const scalar_t* biasptr, \ + scalar_t* cptr, int CuCount, int b_row_stride_bytes_i32, \ + int64_t group_size, int max_lds_len, bool has_zp); \ + \ + template \ + void launch_moe_int4_n##N( \ + dim3 grid, cudaStream_t stream, int K_in, int M_in, int N_in, \ + const uint8_t* wptr, const scalar_t* aptr, const scalar_t* sptr, \ + const scalar_t* zpptr, scalar_t* cptr, const int* eidptr, \ + const int* stidptr, int top_k_in, long expert_stride_w, \ + long expert_stride_s, long expert_stride_zp, int CuCount, \ + int num_expert_blocks, int64_t group_size, int max_lds_len, bool has_zp, \ + bool fuse_silu_mul); + +DECLARE_LAUNCH_INT4(1) +DECLARE_LAUNCH_INT4(2) +DECLARE_LAUNCH_INT4(3) +DECLARE_LAUNCH_INT4(4) +DECLARE_LAUNCH_INT4(5) + +#undef DECLARE_LAUNCH_INT4 diff --git a/csrc/rocm/skinny_gemms_int8.cu b/csrc/rocm/skinny_gemms_int8.cu index bc9abcff031a..ea6d7b7358d3 100644 --- a/csrc/rocm/skinny_gemms_int8.cu +++ b/csrc/rocm/skinny_gemms_int8.cu @@ -12,334 +12,8 @@ #include "../cuda_compat.h" #include "dispatch_utils.h" -#if defined(__HIPCC__) && \ - (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) - #define __HIP__GFX9__ -#endif - -#define LDS_SIZE 64 * 1024 - -int get_lds_size_int8() { - static bool is_cached = false; - static int result; - if (is_cached == false) { - auto dprops = at::cuda::getCurrentDeviceProperties(); - std::string device_arch = dprops->gcnArchName; - size_t substring = device_arch.find("gfx95"); - result = (substring == std::string::npos ? 64 * 1024 : 160 * 1024); - is_cached = true; - } - return result; -} - -bool is_gfx11_int8() { - static const bool result = [] { - auto dprops = at::cuda::getCurrentDeviceProperties(); - std::string device_arch = dprops->gcnArchName; - return device_arch.find("gfx11") != std::string::npos; - }(); - return result; -} - -#if defined(NDEBUG) - #undef NDEBUG - #include - #define UNREACHABLE_CODE assert(false); - #define NDEBUG -#else - #define UNREACHABLE_CODE assert(false); -#endif - -template -__device__ __forceinline__ float __s2float(T v); - -template <> -__device__ __forceinline__ float __s2float(half v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float __s2float(__hip_bfloat16 v) { - return __bfloat162float(v); -} - -template -__device__ __forceinline__ T __float2s(float v); - -template <> -__device__ __forceinline__ half __float2s(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ __hip_bfloat16 __float2s(float v) { - return __float2bfloat16(v); -} - -template -__device__ __forceinline__ T loadnt(T* addr) { - return __builtin_nontemporal_load(addr); -} - -template -struct scalar {}; - -template <> -struct scalar { - using type = half; -}; - -template <> -struct scalar { - using type = __hip_bfloat16; -}; - -#define DOT2C(V0, V2, V3) \ - if constexpr (std::is_same_v) { \ - V0 = __builtin_amdgcn_fdot2(*((half2*)(&(V2))), *((half2*)(&(V3))), V0, \ - false); \ - } else if constexpr (std::is_same_v) { \ - float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ - __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ - V0 += (s.x + s.y); \ - } - -__device__ inline unsigned int min__(uint32_t a, uint32_t b) { - return min(a, b); -} - -// W8A16 skinny GEMM kernel: int8 weights, fp16/bf16 activations -// Targets the "sml" case where activations fit in LDS. -// A_CHUNK=16: each thread processes 16 int8 weight elements per step. -// GROUP_SIZE: 0 = per-channel scale [M] (one scale per output row, applied -// once at the end of the K reduction). -// >0 = per-group scale [M, K/GROUP_SIZE] (one scale per K-group; -// requires GROUP_SIZE % A_CHUNK == 0 so each thread's A_CHUNK -// K-elements lie within a single group). -#if defined(__HIP__GFX9__) || defined(__GFX11__) -template -__global__ void __launch_bounds__(WvPrGrp* THRDS) - wvSplitK_int8_hf_sml_(const int K, const int M, const int Bx, const int By, - const int8_t* B, const scalar_t* __restrict__ A, - const scalar_t* scale, - const scalar_t* __restrict__ BIAS, scalar_t* C, - const int _WvPrGrp, const int CuCount) { - static_assert(GROUP_SIZE == 0 || GROUP_SIZE >= A_CHUNK, - "GROUP_SIZE must be >= A_CHUNK so each thread's A_CHUNK " - "K-elements lie in one group"); - static_assert(GROUP_SIZE == 0 || (GROUP_SIZE % A_CHUNK) == 0, - "GROUP_SIZE must be a multiple of A_CHUNK"); - constexpr int max_lds_len = LDS_SIZE / 2; - - // Activation union: 16 fp16/bf16 values = 32 bytes - union bigTypeA { - scalar_t h[A_CHUNK]; - float f[A_CHUNK / 2]; - }; - - // Weight union: 16 int8 values = 16 bytes - union bigTypeW { - int8_t b[A_CHUNK]; - float f[A_CHUNK / 4]; - }; - - __shared__ scalar_t s[max_lds_len]; - - // Fetch activation matrix to LDS - // Each thread fetches A_CHUNK fp16 elements = 32 bytes - for (uint32_t k = 0; k < min__(K * N, max_lds_len); - k += THRDS * WvPrGrp * A_CHUNK) { - uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); - - if (k_in >= min__(K * N, max_lds_len)) break; - - *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in])); - } - __syncthreads(); - - if (threadIdx.y >= _WvPrGrp) return; - - uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; - - // For per-group scales, num_groups stride along K. - [[maybe_unused]] const int num_groups = - (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0; - - float sum[N][YTILE]; - - while (m < M) { - for (int i = 0; i < YTILE; i++) - for (int n = 0; n < N; n++) sum[n][i] = 0; - - bigTypeA bigA[N][UNRL]; - bigTypeW bigB[YTILE][UNRL]; - - for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { - // Fetch int8 weights from global memory - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - - const int8_t* B_ = &B[(m + 0) * K + k_]; - for (int y = 0; y < YTILE; y++) { - // 16 bytes = 4 floats worth of int8 data - const float* src = (const float*)(&B_[y * K]); - #pragma unroll - for (int i = 0; i < A_CHUNK / 4; i++) - bigB[y][k2].f[i] = loadnt((float*)&src[i]); - } - } - - // Fetch fp16/bf16 activations from LDS - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - - for (int n = 0; n < N; n++) { - bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); - } - } - - // Matrix multiply: convert int8 weight pairs to fp16, then DOT2C - #pragma unroll - for (uint32_t k2 = 0; k2 < UNRL; k2++) { - uint32_t k = k1 + k2 * THRDS * A_CHUNK; - uint32_t k_ = k + threadIdx.x * A_CHUNK; - if (k_ >= K) break; - - #pragma unroll - for (uint32_t n = 0; n < N; n++) { - #pragma unroll - for (int y = 0; y < YTILE; y++) { - // Convert 16 int8 weights to 8 fp16 pairs stored in a bigTypeA - // union - bigTypeA cvtB; - #pragma unroll - for (uint32_t b = 0; b < A_CHUNK; b++) { - cvtB.h[b] = bigB[y][k2].b[b]; - } - if constexpr (GROUP_SIZE > 0) { - // Per-group scale: this thread's A_CHUNK K-elements lie in - // a single group (statically asserted GROUP_SIZE >= A_CHUNK - // and GROUP_SIZE % A_CHUNK == 0). Accumulate the partial - // dot product locally, multiply by the group's scale, then - // add into sum[n][y]. - float partial = 0.0f; - #pragma unroll - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(partial, bigA[n][k2].f[b], cvtB.f[b]) - } - uint32_t group_idx = k_ / GROUP_SIZE; - sum[n][y] += - partial * __s2float(scale[(m + y) * num_groups + group_idx]); - } else { - #pragma unroll - for (uint32_t b = 0; b < A_CHUNK / 2; b++) { - DOT2C(sum[n][y], bigA[n][k2].f[b], cvtB.f[b]) - } - } - } - } - } - } - - // Reduction - #if defined(__GFX11__) - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, - 1); // row_shr8 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, - 1); // row_shr4 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, - 1); // row_shr2 - sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, - 1); // row_shr1 - sum[n][y] += __shfl_xor(sum[n][y], 16); - } - } - - if (threadIdx.x == (THRDS - 1)) { - for (int n = 0; n < N; n++) { - for (int i = 0; i < YTILE; i++) { - if constexpr (GROUP_SIZE == 0) { - sum[n][i] *= __s2float(scale[m + i]); - } - if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]); - C[m + i + n * M] = __float2s(sum[n][i]); - } - } - } - #else // GFX9 wave64 path - for (int n = 0; n < N; n++) { - for (int y = 0; y < YTILE; y++) { - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 " - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 " - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 " - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0" - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0" - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0" - : "=v"(sum[n][y]) - : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); - } - } - - if (threadIdx.x == 63) { - for (int n = 0; n < N; n++) { - for (int i = 0; i < YTILE; i++) { - if constexpr (GROUP_SIZE == 0) { - sum[n][i] *= __s2float(scale[m + i]); - } - if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]); - C[m + i + n * M] = __float2s(sum[n][i]); - } - } - } - #endif // defined(__GFX11__) - m += CuCount * _WvPrGrp * YTILE; - } -} -#else // !defined(__HIP__GFX9__) && !defined(__GFX11__) -template -__global__ void wvSplitK_int8_hf_sml_(const int K, const int M, const int Bx, - const int By, const int8_t* B, - const scalar_t* __restrict__ A, - const scalar_t* scale, - const scalar_t* __restrict__ BIAS, - scalar_t* C, const int _WvPrGrp, - const int CuCount) { - UNREACHABLE_CODE -} -#endif // defined(__HIP__GFX9__) || defined(__GFX11__) - -int mindiv_int8(int N, int div1, int div2) { - int nPrRnd = div1 * div2; - int limit = div2 < 13 ? div2 : 13; - int rnds[16]; - for (int i = 0; i < limit; i++) { - rnds[i] = (N + nPrRnd - 1) / nPrRnd; - nPrRnd -= div1; - } - for (int i = limit - 1; i >= 0; i--) - if (rnds[0] == rnds[i]) return (div2 - i); - return 0; -} +#include "skinny_gemms_int8/kernel.cuh" // arch helpers + scalar<> traits +#include "skinny_gemms_int8/launch.h" // per-N launchers (defined in shards) torch::Tensor wvSplitK_int8(const at::Tensor& in_a, const at::Tensor& in_b, const at::Tensor& in_scale, @@ -403,45 +77,6 @@ torch::Tensor wvSplitK_int8(const at::Tensor& in_a, const at::Tensor& in_b, const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); -#define WVSPLITK_INT8_LAUNCH_G(_THRDS, _YTILE, _UNRL, _N, _GROUP) \ - { \ - dim3 block(_THRDS, 16); \ - int __wvPrGrp = mindiv_int8(M_in, CuCount * _YTILE, 16); \ - TORCH_CHECK(M_in % _YTILE == 0, "M must be divisible by YTILE=", _YTILE); \ - wvSplitK_int8_hf_sml_ \ - <<>>(K_in, M_in, Bx_in, By_in, wptr, aptr, \ - sptr, biasptr, cptr, __wvPrGrp, CuCount); \ - } - -#define WVSPLITK_INT8_LAUNCH(_THRDS, _YTILE, _UNRL, _N) \ - { \ - if (group_size == -1) { \ - WVSPLITK_INT8_LAUNCH_G(_THRDS, _YTILE, _UNRL, _N, 0) \ - } else if (group_size == 32) { \ - WVSPLITK_INT8_LAUNCH_G(_THRDS, _YTILE, _UNRL, _N, 32) \ - } else if (group_size == 64) { \ - WVSPLITK_INT8_LAUNCH_G(_THRDS, _YTILE, _UNRL, _N, 64) \ - } else if (group_size == 128) { \ - WVSPLITK_INT8_LAUNCH_G(_THRDS, _YTILE, _UNRL, _N, 128) \ - } else { \ - TORCH_CHECK(false, "Unsupported group_size=", group_size); \ - } \ - } - -#define WVSPLITK_INT8(_YTILE, _UNRL, _N) \ - if (is_gfx11_int8()) \ - WVSPLITK_INT8_LAUNCH(32, _YTILE, _UNRL, _N) \ - else \ - WVSPLITK_INT8_LAUNCH(64, _YTILE, _UNRL, _N) - -#define WVSPLIT_INT8_TILE(_sYT, __N) \ - { \ - if (__N >= 4 && _sYT >= 480) \ - WVSPLITK_INT8(4, 1, __N) \ - else \ - WVSPLITK_INT8(1, 4, __N) \ - } - AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK_int8", [&] { using fptype = typename scalar::type; const int8_t* wptr = in_a.data_ptr(); @@ -454,22 +89,38 @@ torch::Tensor wvSplitK_int8(const at::Tensor& in_a, const at::Tensor& in_b, fptype* cptr = reinterpret_cast(out_c.data_ptr()); int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); + int thrds = is_gfx11_int8() ? 32 : 64; + int ytile, unrl; + if (N_in >= 4 && sYT >= 480) { + ytile = 4; + unrl = 1; + } else { + ytile = 1; + unrl = 4; + } + + TORCH_CHECK(M_in % ytile == 0, "M must be divisible by YTILE=", ytile); switch (N_in) { case 1: - WVSPLIT_INT8_TILE(sYT, 1) + launch_int8_n1(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, + biasptr, cptr, CuCount, thrds, ytile, unrl, group_size); break; case 2: - WVSPLIT_INT8_TILE(sYT, 2) + launch_int8_n2(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, + biasptr, cptr, CuCount, thrds, ytile, unrl, group_size); break; case 3: - WVSPLIT_INT8_TILE(sYT, 3) + launch_int8_n3(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, + biasptr, cptr, CuCount, thrds, ytile, unrl, group_size); break; case 4: - WVSPLIT_INT8_TILE(sYT, 4) + launch_int8_n4(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, + biasptr, cptr, CuCount, thrds, ytile, unrl, group_size); break; case 5: - WVSPLIT_INT8_TILE(sYT, 5) + launch_int8_n5(grid, stream, K_in, M_in, Bx_in, By_in, wptr, aptr, sptr, + biasptr, cptr, CuCount, thrds, ytile, unrl, group_size); break; default: throw std::runtime_error( @@ -478,11 +129,6 @@ torch::Tensor wvSplitK_int8(const at::Tensor& in_a, const at::Tensor& in_b, } }); -#undef WVSPLITK_INT8_LAUNCH_G -#undef WVSPLITK_INT8_LAUNCH -#undef WVSPLITK_INT8 -#undef WVSPLIT_INT8_TILE - return out_c; } @@ -531,91 +177,31 @@ torch::Tensor wvSplitK_int8_sweep(const at::Tensor& in_a, const int THRDS = is_gfx11_int8() ? 32 : 64; - #define SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL, _N) \ - { \ - dim3 block(_THRDS, _WVPRGRP); \ - int __wvPrGrp = mindiv_int8(M_in, CuCount * _YTILE, _WVPRGRP); \ - wvSplitK_int8_hf_sml_ \ - <<>>(K_in, M_in, 1, 1, wptr, aptr, sptr, \ - biasptr, cptr, __wvPrGrp, CuCount); \ - } - - #define SWEEP_N(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL) \ - switch (N_in) { \ - case 1: \ - SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL, 1) break; \ - case 2: \ - SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL, 2) break; \ - case 3: \ - SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL, 3) break; \ - case 4: \ - SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL, 4) break; \ - default: \ - TORCH_CHECK(false, "Unsupported N=", N_in); \ - } - - #define SWEEP_UNRL(_THRDS, _YTILE, _WVPRGRP, _ACHUNK) \ - if (unrl == 1) { \ - SWEEP_N(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 1) \ - } else if (unrl == 2) { \ - SWEEP_N(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 2) \ - } else if (unrl == 4) { \ - SWEEP_N(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 4) \ - } else { \ - TORCH_CHECK(false, "Unsupported unrl=", unrl); \ - } - - #define SWEEP_YTILE(_THRDS, _WVPRGRP, _ACHUNK) \ - if (ytile == 1) { \ - SWEEP_UNRL(_THRDS, 1, _WVPRGRP, _ACHUNK) \ - } else if (ytile == 2) { \ - SWEEP_UNRL(_THRDS, 2, _WVPRGRP, _ACHUNK) \ - } else if (ytile == 4) { \ - SWEEP_UNRL(_THRDS, 4, _WVPRGRP, _ACHUNK) \ - } else { \ - TORCH_CHECK(false, "Unsupported ytile=", ytile); \ - } - - #define SWEEP_WVPRGRP(_THRDS, _ACHUNK) \ - if (wvprgrp == 8) { \ - SWEEP_YTILE(_THRDS, 8, _ACHUNK) \ - } else if (wvprgrp == 12) { \ - SWEEP_YTILE(_THRDS, 12, _ACHUNK) \ - } else if (wvprgrp == 16) { \ - SWEEP_YTILE(_THRDS, 16, _ACHUNK) \ - } else { \ - TORCH_CHECK(false, "Unsupported wvprgrp=", wvprgrp); \ - } - - if (THRDS == 32) { - if (achunk == 8) { - SWEEP_WVPRGRP(32, 8) - } else if (achunk == 16) { - SWEEP_WVPRGRP(32, 16) - } else if (achunk == 32) { - SWEEP_WVPRGRP(32, 32) - } else { - TORCH_CHECK(false, "Unsupported achunk=", achunk); - } - } else { - if (achunk == 8) { - SWEEP_WVPRGRP(64, 8) - } else if (achunk == 16) { - SWEEP_WVPRGRP(64, 16) - } else if (achunk == 32) { - SWEEP_WVPRGRP(64, 32) - } else { - TORCH_CHECK(false, "Unsupported achunk=", achunk); - } + switch (N_in) { + case 1: + launch_int8_n1_sweep(grid, stream, K_in, M_in, 1, 1, wptr, aptr, sptr, + biasptr, cptr, CuCount, THRDS, ytile, wvprgrp, + achunk, unrl); + break; + case 2: + launch_int8_n2_sweep(grid, stream, K_in, M_in, 1, 1, wptr, aptr, sptr, + biasptr, cptr, CuCount, THRDS, ytile, wvprgrp, + achunk, unrl); + break; + case 3: + launch_int8_n3_sweep(grid, stream, K_in, M_in, 1, 1, wptr, aptr, sptr, + biasptr, cptr, CuCount, THRDS, ytile, wvprgrp, + achunk, unrl); + break; + case 4: + launch_int8_n4_sweep(grid, stream, K_in, M_in, 1, 1, wptr, aptr, sptr, + biasptr, cptr, CuCount, THRDS, ytile, wvprgrp, + achunk, unrl); + break; + default: + TORCH_CHECK(false, "Unsupported N=", N_in); } - #undef SWEEP_LAUNCH - #undef SWEEP_N - #undef SWEEP_UNRL - #undef SWEEP_YTILE - #undef SWEEP_WVPRGRP - return out_c; } #endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/dispatch.cuh b/csrc/rocm/skinny_gemms_int8/dispatch.cuh new file mode 100644 index 000000000000..d5200afadaa1 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/dispatch.cuh @@ -0,0 +1,117 @@ +#pragma once + +#include "kernel.cuh" + +// Production launch dispatch for wvSplitK_int8. +// The heuristic in the parent .cu picks (ytile, unrl) and group_size, +// then calls into the per-N shard via launch_int8_nX(). +template +inline void dispatch_int8(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int unrl, int64_t group_size) { +#define LAUNCH_G(_THRDS, _YTILE, _UNRL, _GROUP) \ + do { \ + dim3 block(_THRDS, 16); \ + int __wvPrGrp = mindiv_int8(M, CuCount * (_YTILE), 16); \ + wvSplitK_int8_hf_sml_<<>>( \ + K, M, Bx, By, B, A, scale, BIAS, C, __wvPrGrp, CuCount); \ + return; \ + } while (0) + +#define LAUNCH(_THRDS, _YTILE, _UNRL) \ + do { \ + if (group_size == -1) { \ + LAUNCH_G(_THRDS, _YTILE, _UNRL, 0); \ + } else if (group_size == 32) { \ + LAUNCH_G(_THRDS, _YTILE, _UNRL, 32); \ + } else if (group_size == 64) { \ + LAUNCH_G(_THRDS, _YTILE, _UNRL, 64); \ + } else if (group_size == 128) { \ + LAUNCH_G(_THRDS, _YTILE, _UNRL, 128); \ + } \ + } while (0) + + // Enumerate the (thrds, ytile, unrl) tuples WVSPLIT_INT8_TILE picks. + if (thrds == 32) { + if (ytile == 4 && unrl == 1) LAUNCH(32, 4, 1); + if (ytile == 1 && unrl == 4) LAUNCH(32, 1, 4); + } else { + if (ytile == 4 && unrl == 1) LAUNCH(64, 4, 1); + if (ytile == 1 && unrl == 4) LAUNCH(64, 1, 4); + } + TORCH_CHECK(false, "wvSplitK_int8: unhandled (thrds=", thrds, + ", ytile=", ytile, ", unrl=", unrl, ", group_size=", group_size, + "). Add to dispatch.cuh."); + +#undef LAUNCH_G +#undef LAUNCH +} + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +// Sweep dispatch: full (yt, ur, ac, wv) cross-product. +template +inline void dispatch_int8_sweep(dim3 grid, cudaStream_t stream, int K, int M, + int Bx, int By, const int8_t* B, + const scalar_t* A, const scalar_t* scale, + const scalar_t* BIAS, scalar_t* C, int CuCount, + int thrds, int ytile, int wvprgrp, int achunk, + int unrl) { + #define SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, _UNRL) \ + do { \ + dim3 block(_THRDS, _WVPRGRP); \ + int __wvPrGrp = mindiv_int8(M, CuCount * (_YTILE), _WVPRGRP); \ + wvSplitK_int8_hf_sml_<<>>( \ + K, M, Bx, By, B, A, scale, BIAS, C, __wvPrGrp, CuCount); \ + return; \ + } while (0) + + #define SWEEP_UNRL(_THRDS, _YTILE, _WVPRGRP, _ACHUNK) \ + do { \ + if (unrl == 1) SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 1); \ + if (unrl == 2) SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 2); \ + if (unrl == 4) SWEEP_LAUNCH(_THRDS, _YTILE, _WVPRGRP, _ACHUNK, 4); \ + } while (0) + + #define SWEEP_YTILE(_THRDS, _WVPRGRP, _ACHUNK) \ + do { \ + if (ytile == 1) SWEEP_UNRL(_THRDS, 1, _WVPRGRP, _ACHUNK); \ + if (ytile == 2) SWEEP_UNRL(_THRDS, 2, _WVPRGRP, _ACHUNK); \ + if (ytile == 4) SWEEP_UNRL(_THRDS, 4, _WVPRGRP, _ACHUNK); \ + } while (0) + + #define SWEEP_WV(_THRDS, _ACHUNK) \ + do { \ + if (wvprgrp == 8) SWEEP_YTILE(_THRDS, 8, _ACHUNK); \ + if (wvprgrp == 12) SWEEP_YTILE(_THRDS, 12, _ACHUNK); \ + if (wvprgrp == 16) SWEEP_YTILE(_THRDS, 16, _ACHUNK); \ + } while (0) + + #define SWEEP_AC(_THRDS) \ + do { \ + if (achunk == 8) SWEEP_WV(_THRDS, 8); \ + if (achunk == 16) SWEEP_WV(_THRDS, 16); \ + if (achunk == 32) SWEEP_WV(_THRDS, 32); \ + } while (0) + + if (thrds == 32) { + SWEEP_AC(32); + } else { + SWEEP_AC(64); + } + TORCH_CHECK(false, "wvSplitK_int8_sweep: unhandled (thrds=", thrds, + ", ytile=", ytile, ", wvprgrp=", wvprgrp, ", achunk=", achunk, + ", unrl=", unrl, ")"); + + #undef SWEEP_LAUNCH + #undef SWEEP_UNRL + #undef SWEEP_YTILE + #undef SWEEP_WV + #undef SWEEP_AC +} + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/instantiate_n1.cu b/csrc/rocm/skinny_gemms_int8/instantiate_n1.cu new file mode 100644 index 000000000000..1cc9033a46e8 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/instantiate_n1.cu @@ -0,0 +1,49 @@ +// Per-N kernel instantiation shard for wvSplitK_int8 (N=1). +// One TU per N value keeps the template instantiation footprint per file +// small and lets make parallelize across the shards. See dispatch.cuh. + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int8_n1(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, scalar_t* C, + int CuCount, int thrds, int ytile, int unrl, + int64_t group_size) { + dispatch_int8(grid, stream, K, M, Bx, By, B, A, scale, BIAS, C, + CuCount, thrds, ytile, unrl, group_size); +} + +template void launch_int8_n1(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, const half*, + const half*, half*, int, int, int, int, + int64_t); +template void launch_int8_n1<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int64_t); + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +template +void launch_int8_n1_sweep(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int wvprgrp, int achunk, int unrl) { + dispatch_int8_sweep(grid, stream, K, M, Bx, By, B, A, scale, + BIAS, C, CuCount, thrds, ytile, wvprgrp, + achunk, unrl); +} + +template void launch_int8_n1_sweep(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, + const half*, const half*, half*, int, + int, int, int, int, int); +template void launch_int8_n1_sweep<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int, int); + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/instantiate_n2.cu b/csrc/rocm/skinny_gemms_int8/instantiate_n2.cu new file mode 100644 index 000000000000..3c85ba2d38c9 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/instantiate_n2.cu @@ -0,0 +1,47 @@ +// Per-N kernel instantiation shard for wvSplitK_int8 (N=2). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int8_n2(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, scalar_t* C, + int CuCount, int thrds, int ytile, int unrl, + int64_t group_size) { + dispatch_int8(grid, stream, K, M, Bx, By, B, A, scale, BIAS, C, + CuCount, thrds, ytile, unrl, group_size); +} + +template void launch_int8_n2(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, const half*, + const half*, half*, int, int, int, int, + int64_t); +template void launch_int8_n2<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int64_t); + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +template +void launch_int8_n2_sweep(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int wvprgrp, int achunk, int unrl) { + dispatch_int8_sweep(grid, stream, K, M, Bx, By, B, A, scale, + BIAS, C, CuCount, thrds, ytile, wvprgrp, + achunk, unrl); +} + +template void launch_int8_n2_sweep(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, + const half*, const half*, half*, int, + int, int, int, int, int); +template void launch_int8_n2_sweep<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int, int); + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/instantiate_n3.cu b/csrc/rocm/skinny_gemms_int8/instantiate_n3.cu new file mode 100644 index 000000000000..dccc34f14190 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/instantiate_n3.cu @@ -0,0 +1,47 @@ +// Per-N kernel instantiation shard for wvSplitK_int8 (N=3). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int8_n3(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, scalar_t* C, + int CuCount, int thrds, int ytile, int unrl, + int64_t group_size) { + dispatch_int8(grid, stream, K, M, Bx, By, B, A, scale, BIAS, C, + CuCount, thrds, ytile, unrl, group_size); +} + +template void launch_int8_n3(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, const half*, + const half*, half*, int, int, int, int, + int64_t); +template void launch_int8_n3<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int64_t); + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +template +void launch_int8_n3_sweep(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int wvprgrp, int achunk, int unrl) { + dispatch_int8_sweep(grid, stream, K, M, Bx, By, B, A, scale, + BIAS, C, CuCount, thrds, ytile, wvprgrp, + achunk, unrl); +} + +template void launch_int8_n3_sweep(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, + const half*, const half*, half*, int, + int, int, int, int, int); +template void launch_int8_n3_sweep<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int, int); + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/instantiate_n4.cu b/csrc/rocm/skinny_gemms_int8/instantiate_n4.cu new file mode 100644 index 000000000000..3d4026fd7e5f --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/instantiate_n4.cu @@ -0,0 +1,47 @@ +// Per-N kernel instantiation shard for wvSplitK_int8 (N=4). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int8_n4(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, scalar_t* C, + int CuCount, int thrds, int ytile, int unrl, + int64_t group_size) { + dispatch_int8(grid, stream, K, M, Bx, By, B, A, scale, BIAS, C, + CuCount, thrds, ytile, unrl, group_size); +} + +template void launch_int8_n4(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, const half*, + const half*, half*, int, int, int, int, + int64_t); +template void launch_int8_n4<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int64_t); + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +template +void launch_int8_n4_sweep(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int wvprgrp, int achunk, int unrl) { + dispatch_int8_sweep(grid, stream, K, M, Bx, By, B, A, scale, + BIAS, C, CuCount, thrds, ytile, wvprgrp, + achunk, unrl); +} + +template void launch_int8_n4_sweep(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, + const half*, const half*, half*, int, + int, int, int, int, int); +template void launch_int8_n4_sweep<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int, int); + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/instantiate_n5.cu b/csrc/rocm/skinny_gemms_int8/instantiate_n5.cu new file mode 100644 index 000000000000..e2254b06d07e --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/instantiate_n5.cu @@ -0,0 +1,47 @@ +// Per-N kernel instantiation shard for wvSplitK_int8 (N=5). + +#include "dispatch.cuh" +#include "launch.h" + +template +void launch_int8_n5(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, scalar_t* C, + int CuCount, int thrds, int ytile, int unrl, + int64_t group_size) { + dispatch_int8(grid, stream, K, M, Bx, By, B, A, scale, BIAS, C, + CuCount, thrds, ytile, unrl, group_size); +} + +template void launch_int8_n5(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, const half*, + const half*, half*, int, int, int, int, + int64_t); +template void launch_int8_n5<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int64_t); + +#ifdef VLLM_SKINNY_GEMM_SWEEP + +template +void launch_int8_n5_sweep(dim3 grid, cudaStream_t stream, int K, int M, int Bx, + int By, const int8_t* B, const scalar_t* A, + const scalar_t* scale, const scalar_t* BIAS, + scalar_t* C, int CuCount, int thrds, int ytile, + int wvprgrp, int achunk, int unrl) { + dispatch_int8_sweep(grid, stream, K, M, Bx, By, B, A, scale, + BIAS, C, CuCount, thrds, ytile, wvprgrp, + achunk, unrl); +} + +template void launch_int8_n5_sweep(dim3, cudaStream_t, int, int, int, int, + const int8_t*, const half*, + const half*, const half*, half*, int, + int, int, int, int, int); +template void launch_int8_n5_sweep<__hip_bfloat16>( + dim3, cudaStream_t, int, int, int, int, const int8_t*, + const __hip_bfloat16*, const __hip_bfloat16*, const __hip_bfloat16*, + __hip_bfloat16*, int, int, int, int, int, int); + +#endif // VLLM_SKINNY_GEMM_SWEEP diff --git a/csrc/rocm/skinny_gemms_int8/kernel.cuh b/csrc/rocm/skinny_gemms_int8/kernel.cuh new file mode 100644 index 000000000000..04b466661e54 --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/kernel.cuh @@ -0,0 +1,342 @@ +#pragma once + +#include + +#include +#include +#include + +#include +#include +#include + +#include "../../cuda_compat.h" + +#if defined(__HIPCC__) && \ + (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__GFX9__ +#endif + +#define LDS_SIZE 64 * 1024 + +inline int get_lds_size_int8() { + static bool is_cached = false; + static int result; + if (is_cached == false) { + auto dprops = at::cuda::getCurrentDeviceProperties(); + std::string device_arch = dprops->gcnArchName; + size_t substring = device_arch.find("gfx95"); + result = (substring == std::string::npos ? 64 * 1024 : 160 * 1024); + is_cached = true; + } + return result; +} + +inline bool is_gfx11_int8() { + static const bool result = [] { + auto dprops = at::cuda::getCurrentDeviceProperties(); + std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx11") != std::string::npos; + }(); + return result; +} + +#if defined(NDEBUG) + #undef NDEBUG + #include + #define UNREACHABLE_CODE assert(false); + #define NDEBUG +#else + #define UNREACHABLE_CODE assert(false); +#endif + +template +__device__ __forceinline__ float __s2float(T v); + +template <> +__device__ __forceinline__ float __s2float(half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float __s2float(__hip_bfloat16 v) { + return __bfloat162float(v); +} + +template +__device__ __forceinline__ T __float2s(float v); + +template <> +__device__ __forceinline__ half __float2s(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ __hip_bfloat16 __float2s(float v) { + return __float2bfloat16(v); +} + +template +__device__ __forceinline__ T loadnt(T* addr) { + return __builtin_nontemporal_load(addr); +} + +template +struct scalar {}; + +template <> +struct scalar { + using type = half; +}; + +template <> +struct scalar { + using type = __hip_bfloat16; +}; + +#define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + V0 = __builtin_amdgcn_fdot2(*((half2*)(&(V2))), *((half2*)(&(V3))), V0, \ + false); \ + } else if constexpr (std::is_same_v) { \ + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ + V0 += (s.x + s.y); \ + } + +__device__ inline unsigned int min__(uint32_t a, uint32_t b) { + return min(a, b); +} + +inline int mindiv_int8(int N, int div1, int div2) { + int nPrRnd = div1 * div2; + int limit = div2 < 13 ? div2 : 13; + int rnds[16]; + for (int i = 0; i < limit; i++) { + rnds[i] = (N + nPrRnd - 1) / nPrRnd; + nPrRnd -= div1; + } + for (int i = limit - 1; i >= 0; i--) + if (rnds[0] == rnds[i]) return (div2 - i); + return 0; +} + +// W8A16 skinny GEMM kernel: int8 weights, fp16/bf16 activations +// Targets the "sml" case where activations fit in LDS. +// A_CHUNK=16: each thread processes 16 int8 weight elements per step. +// GROUP_SIZE: 0 = per-channel scale [M] (one scale per output row, applied +// once at the end of the K reduction). +// >0 = per-group scale [M, K/GROUP_SIZE] (one scale per K-group; +// requires GROUP_SIZE % A_CHUNK == 0 so each thread's A_CHUNK +// K-elements lie within a single group). +#if defined(__HIP__GFX9__) || defined(__GFX11__) +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_int8_hf_sml_(const int K, const int M, const int Bx, const int By, + const int8_t* B, const scalar_t* __restrict__ A, + const scalar_t* scale, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + static_assert(GROUP_SIZE == 0 || GROUP_SIZE >= A_CHUNK, + "GROUP_SIZE must be >= A_CHUNK so each thread's A_CHUNK " + "K-elements lie in one group"); + static_assert(GROUP_SIZE == 0 || (GROUP_SIZE % A_CHUNK) == 0, + "GROUP_SIZE must be a multiple of A_CHUNK"); + constexpr int max_lds_len = LDS_SIZE / 2; + + // Activation union: 16 fp16/bf16 values = 32 bytes + union bigTypeA { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + }; + + // Weight union: 16 int8 values = 16 bytes + union bigTypeW { + int8_t b[A_CHUNK]; + float f[A_CHUNK / 4]; + }; + + __shared__ scalar_t s[max_lds_len]; + + // Fetch activation matrix to LDS + // Each thread fetches A_CHUNK fp16 elements = 32 bytes + for (uint32_t k = 0; k < min__(K * N, max_lds_len); + k += THRDS * WvPrGrp * A_CHUNK) { + uint32_t k_in = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); + + if (k_in >= min__(K * N, max_lds_len)) break; + + *((bigTypeA*)(&s[k_in])) = *((bigTypeA*)(&A[k_in])); + } + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; + + // For per-group scales, num_groups stride along K. + [[maybe_unused]] const int num_groups = + (GROUP_SIZE > 0) ? (K / GROUP_SIZE) : 0; + + float sum[N][YTILE]; + + while (m < M) { + for (int i = 0; i < YTILE; i++) + for (int n = 0; n < N; n++) sum[n][i] = 0; + + bigTypeA bigA[N][UNRL]; + bigTypeW bigB[YTILE][UNRL]; + + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + // Fetch int8 weights from global memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + + const int8_t* B_ = &B[(m + 0) * K + k_]; + for (int y = 0; y < YTILE; y++) { + // 16 bytes = 4 floats worth of int8 data + const float* src = (const float*)(&B_[y * K]); + #pragma unroll + for (int i = 0; i < A_CHUNK / 4; i++) + bigB[y][k2].f[i] = loadnt((float*)&src[i]); + } + } + + // Fetch fp16/bf16 activations from LDS + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigTypeA*)(&(s[k_ + K * n]))); + } + } + + // Matrix multiply: convert int8 weight pairs to fp16, then DOT2C + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + + #pragma unroll + for (uint32_t n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + // Convert 16 int8 weights to 8 fp16 pairs stored in a bigTypeA + // union + bigTypeA cvtB; + #pragma unroll + for (uint32_t b = 0; b < A_CHUNK; b++) { + cvtB.h[b] = bigB[y][k2].b[b]; + } + if constexpr (GROUP_SIZE > 0) { + // Per-group scale: this thread's A_CHUNK K-elements lie in + // a single group (statically asserted GROUP_SIZE >= A_CHUNK + // and GROUP_SIZE % A_CHUNK == 0). Accumulate the partial + // dot product locally, multiply by the group's scale, then + // add into sum[n][y]. + float partial = 0.0f; + #pragma unroll + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(partial, bigA[n][k2].f[b], cvtB.f[b]) + } + uint32_t group_idx = k_ / GROUP_SIZE; + sum[n][y] += + partial * __s2float(scale[(m + y) * num_groups + group_idx]); + } else { + #pragma unroll + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], cvtB.f[b]) + } + } + } + } + } + } + + // Reduction + #if defined(__GFX11__) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + sum[n][y] += __shfl_xor(sum[n][y], 16); + } + } + + if (threadIdx.x == (THRDS - 1)) { + for (int n = 0; n < N; n++) { + for (int i = 0; i < YTILE; i++) { + if constexpr (GROUP_SIZE == 0) { + sum[n][i] *= __s2float(scale[m + i]); + } + if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]); + C[m + i + n * M] = __float2s(sum[n][i]); + } + } + } + #else // GFX9 wave64 path + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 wave_shr:1 bound_ctrl:0" + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:15 bound_ctrl:0" + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_bcast:31 bound_ctrl:0" + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + } + } + + if (threadIdx.x == 63) { + for (int n = 0; n < N; n++) { + for (int i = 0; i < YTILE; i++) { + if constexpr (GROUP_SIZE == 0) { + sum[n][i] *= __s2float(scale[m + i]); + } + if (BIAS) sum[n][i] += __s2float(BIAS[(m + i) % Bx + (n % By) * M]); + C[m + i + n * M] = __float2s(sum[n][i]); + } + } + } + #endif // defined(__GFX11__) + m += CuCount * _WvPrGrp * YTILE; + } +} +#else // !defined(__HIP__GFX9__) && !defined(__GFX11__) +template +__global__ void wvSplitK_int8_hf_sml_(const int K, const int M, const int Bx, + const int By, const int8_t* B, + const scalar_t* __restrict__ A, + const scalar_t* scale, + const scalar_t* __restrict__ BIAS, + scalar_t* C, const int _WvPrGrp, + const int CuCount) { + UNREACHABLE_CODE +} +#endif // defined(__HIP__GFX9__) || defined(__GFX11__) diff --git a/csrc/rocm/skinny_gemms_int8/launch.h b/csrc/rocm/skinny_gemms_int8/launch.h new file mode 100644 index 000000000000..82d36a4177ef --- /dev/null +++ b/csrc/rocm/skinny_gemms_int8/launch.h @@ -0,0 +1,49 @@ +#pragma once + +#include +#include +#include + +#include + +// Per-N launchers for the wvSplitK_int8 kernel. Each instantiate_n{K}.cu TU +// defines launch_int8_n{K}() for a single N value, which keeps the template +// instantiation footprint per TU small and lets make parallelize across the +// shards. The dispatcher in the parent .cu picks (yt, ur) via the heuristic +// and then calls into the appropriate shard. + +#define DECLARE_LAUNCH_INT8(N) \ + template \ + void launch_int8_n##N(dim3 grid, cudaStream_t stream, int K, int M, int Bx, \ + int By, const int8_t* B, const scalar_t* A, \ + const scalar_t* scale, const scalar_t* BIAS, \ + scalar_t* C, int CuCount, int thrds, int ytile, \ + int unrl, int64_t group_size); + +DECLARE_LAUNCH_INT8(1) +DECLARE_LAUNCH_INT8(2) +DECLARE_LAUNCH_INT8(3) +DECLARE_LAUNCH_INT8(4) +DECLARE_LAUNCH_INT8(5) + +#undef DECLARE_LAUNCH_INT8 + +#ifdef VLLM_SKINNY_GEMM_SWEEP + + #define DECLARE_LAUNCH_INT8_SWEEP(N) \ + template \ + void launch_int8_n##N##_sweep( \ + dim3 grid, cudaStream_t stream, int K, int M, int Bx, int By, \ + const int8_t* B, const scalar_t* A, const scalar_t* scale, \ + const scalar_t* BIAS, scalar_t* C, int CuCount, int thrds, int ytile, \ + int wvprgrp, int achunk, int unrl); + +DECLARE_LAUNCH_INT8_SWEEP(1) +DECLARE_LAUNCH_INT8_SWEEP(2) +DECLARE_LAUNCH_INT8_SWEEP(3) +DECLARE_LAUNCH_INT8_SWEEP(4) +DECLARE_LAUNCH_INT8_SWEEP(5) + + #undef DECLARE_LAUNCH_INT8_SWEEP + +#endif // VLLM_SKINNY_GEMM_SWEEP